vyre_driver_cuda/
megakernel_plan_cache.rs

1//! Bounded CUDA megakernel plan cache.
2//!
3//! The cache stores topology decisions keyed by stable graph layout,
4//! analysis family, CUDA device feature signature, and coarse runtime-pressure
5//! buckets. The first three fields are the architectural identity of a plan;
6//! pressure buckets prevent a sparse first query from poisoning dense later
7//! queries over the same resident graph.
8
9use std::cmp::Reverse;
10use std::collections::BinaryHeap;
11
12use rustc_hash::FxHashMap;
13
14use crate::backend::ordering::sort_unstable_by_key_if_needed;
15use crate::backend::staging_reserve::reserve_vec;
16use crate::device::CudaDeviceCaps;
17use crate::megakernel_scheduler::{
18    select_cuda_megakernel_topology, select_cuda_megakernel_topology_stable,
19    CudaMegakernelScheduleSample,
20};
21use vyre_driver::megakernel_execution::{
22    plan_megakernel_memory_budget, MegakernelExecutionPlan, MegakernelExecutionTopology,
23    MegakernelGraphShape, MegakernelMemoryBudget, MegakernelMemoryError,
24    MegakernelTopologyDecision,
25};
26
27const DEFAULT_MAX_MEGAKERNEL_PLANS: usize = 256;
28const PRESSURE_BUCKET_BPS: u32 = 1_000;
29const DENSITY_BUCKETS: u16 = 16;
30const READBACK_BUCKET_SHIFT: u32 = 12;
31
32/// Analysis family for a cached CUDA megakernel plan.
33#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
34pub enum CudaMegakernelAnalysisKind {
35    /// Generic graph dataflow wave.
36    Dataflow,
37    /// IFDS/IDE-style exploded-supergraph propagation.
38    Ifds,
39    /// Reaching-definitions propagation.
40    ReachingDefinitions,
41    /// Live-variable propagation.
42    Liveness,
43    /// Points-to propagation.
44    PointsTo,
45    /// Source-token or parser-frontier wave.
46    ParserFrontend,
47    /// Caller-owned analysis family identified by a stable numeric tag.
48    Custom(u64),
49}
50
51/// CUDA device feature signature that invalidates cached megakernel plans.
52#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
53pub struct CudaMegakernelDeviceKey {
54    /// CUDA SM major version.
55    pub sm_major: u16,
56    /// CUDA SM minor version.
57    pub sm_minor: u16,
58    /// Hardware warp size.
59    pub warp_size: u16,
60    /// Whether cooperative grid synchronization is available.
61    pub supports_grid_sync: bool,
62    /// Whether tensor-core lowering is available for this backend session.
63    pub supports_tensor_cores: bool,
64    /// Maximum threads accepted for one workgroup/block.
65    pub max_workgroup_size: u32,
66}
67
68impl From<&CudaDeviceCaps> for CudaMegakernelDeviceKey {
69    fn from(caps: &CudaDeviceCaps) -> Self {
70        Self {
71            sm_major: caps.compute_capability.0.min(u32::from(u16::MAX)) as u16,
72            sm_minor: caps.compute_capability.1.min(u32::from(u16::MAX)) as u16,
73            warp_size: caps.required_warp_size_u32().min(u32::from(u16::MAX)) as u16,
74            supports_grid_sync: caps.compute_capability >= (6, 0) && caps.cooperative_launch,
75            supports_tensor_cores: caps.hardware_supports_tensor_cores(),
76            max_workgroup_size: caps.max_threads_per_block_u32(),
77        }
78    }
79}
80
81/// Stable key for cached CUDA megakernel plans.
82#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
83pub struct CudaMegakernelPlanCacheKey {
84    /// Stable hash of the normalized resident graph layout.
85    pub graph_layout_hash: u64,
86    /// Analysis family consuming the graph layout.
87    pub analysis_kind: CudaMegakernelAnalysisKind,
88    /// CUDA device feature signature.
89    pub device: CudaMegakernelDeviceKey,
90    /// Coarse active-frontier density bucket.
91    pub frontier_density_bucket: u16,
92    /// Coarse memory-pressure bucket in basis points.
93    pub memory_pressure_bucket: u32,
94    /// Coarse output/readback pressure bucket.
95    pub readback_pressure_bucket: u16,
96    /// Coarse launch-over-dispatch pressure bucket in basis points.
97    pub launch_pressure_bucket: u32,
98    /// Coarse caller-provided fusion-pressure bucket.
99    pub fusion_pressure_bucket: u32,
100}
101
102impl CudaMegakernelPlanCacheKey {
103    /// Build a cache key from stable identity fields and runtime pressure.
104    #[must_use]
105    pub fn new(
106        graph_layout_hash: u64,
107        analysis_kind: CudaMegakernelAnalysisKind,
108        device: CudaMegakernelDeviceKey,
109        frontier_density: f64,
110        memory_pressure_bps: u32,
111        readback_bytes: u64,
112        launch_pressure_bps: u32,
113        fusion_pressure: f64,
114    ) -> Self {
115        Self {
116            graph_layout_hash,
117            analysis_kind,
118            device,
119            frontier_density_bucket: density_bucket(frontier_density),
120            memory_pressure_bucket: pressure_bucket(memory_pressure_bps),
121            readback_pressure_bucket: readback_bucket(readback_bytes),
122            launch_pressure_bucket: pressure_bucket(launch_pressure_bps),
123            fusion_pressure_bucket: fusion_bucket(fusion_pressure),
124        }
125    }
126
127    fn identity(self) -> CudaMegakernelPlanIdentityKey {
128        CudaMegakernelPlanIdentityKey {
129            graph_layout_hash: self.graph_layout_hash,
130            analysis_kind: self.analysis_kind,
131            device: self.device,
132        }
133    }
134}
135
136#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
137struct CudaMegakernelPlanIdentityKey {
138    graph_layout_hash: u64,
139    analysis_kind: CudaMegakernelAnalysisKind,
140    device: CudaMegakernelDeviceKey,
141}
142
143/// Cached CUDA megakernel plan.
144#[derive(Clone, Copy, Debug, PartialEq)]
145pub struct CudaMegakernelCachedPlan {
146    /// Selected topology for this key.
147    pub topology: MegakernelExecutionTopology,
148    /// Full decision telemetry used when the plan was inserted.
149    pub decision: MegakernelTopologyDecision,
150}
151
152/// Runtime counters for [`CudaMegakernelPlanCache`].
153#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
154pub struct CudaMegakernelPlanCacheStats {
155    /// Cache lookup hits.
156    pub hits: u64,
157    /// Cache lookup misses.
158    pub misses: u64,
159    /// Entries evicted by the bounded LRU policy.
160    pub evictions: u64,
161    /// Current entry count.
162    pub entries: usize,
163}
164
165#[derive(Clone, Copy, Debug)]
166struct CudaMegakernelPlanCacheEntry {
167    plan: CudaMegakernelCachedPlan,
168    last_seen: u64,
169}
170
171/// Bounded LRU cache for CUDA megakernel topology plans.
172#[derive(Debug)]
173pub struct CudaMegakernelPlanCache {
174    entries: FxHashMap<CudaMegakernelPlanCacheKey, CudaMegakernelPlanCacheEntry>,
175    latest_by_identity:
176        FxHashMap<CudaMegakernelPlanIdentityKey, (u64, MegakernelExecutionTopology)>,
177    eviction_queue: BinaryHeap<Reverse<(u64, CudaMegakernelPlanCacheKey)>>,
178    max_entries: usize,
179    serial: u64,
180    hits: u64,
181    misses: u64,
182    evictions: u64,
183}
184
185fn increment_plan_cache_counter(counter: &mut u64, field: &'static str) {
186    vyre_driver::accounting::pinning_increment_u64(counter, || {
187        tracing::error!(
188            "CUDA megakernel {field} overflowed u64; pinning counter at u64::MAX. Fix: scrape metrics more frequently or shard the cache."
189        );
190    });
191}
192
193impl Default for CudaMegakernelPlanCache {
194    fn default() -> Self {
195        Self::new()
196    }
197}
198
199impl CudaMegakernelPlanCache {
200    /// Create a cache with the default production entry bound.
201    #[must_use]
202    pub fn new() -> Self {
203        Self::with_max_entries(DEFAULT_MAX_MEGAKERNEL_PLANS)
204    }
205
206    /// Create a cache with an explicit entry bound.
207    #[must_use]
208    pub fn with_max_entries(max_entries: usize) -> Self {
209        Self {
210            entries: FxHashMap::default(),
211            latest_by_identity: FxHashMap::default(),
212            eviction_queue: BinaryHeap::new(),
213            max_entries,
214            serial: 0,
215            hits: 0,
216            misses: 0,
217            evictions: 0,
218        }
219    }
220
221    /// Return a cached plan or insert a newly selected topology decision.
222    pub fn get_or_insert_with(
223        &mut self,
224        key: CudaMegakernelPlanCacheKey,
225        build: impl FnOnce() -> MegakernelTopologyDecision,
226    ) -> Result<CudaMegakernelCachedPlan, MegakernelMemoryError> {
227        let serial = self.advance_serial()?;
228        if let Some(entry) = self.entries.get_mut(&key) {
229            increment_plan_cache_counter(&mut self.hits, "megakernel plan-cache hit counter");
230            entry.last_seen = serial;
231            let plan = entry.plan;
232            self.eviction_queue.push(Reverse((serial, key)));
233            self.update_latest_identity(key.identity(), serial, plan.topology);
234            return Ok(plan);
235        }
236        increment_plan_cache_counter(&mut self.misses, "megakernel plan-cache miss counter");
237        if self.max_entries == 0 {
238            let decision = build();
239            return Ok(CudaMegakernelCachedPlan {
240                topology: decision.topology,
241                decision,
242            });
243        }
244        self.evict_until_below_limit()?;
245        let decision = build();
246        let plan = CudaMegakernelCachedPlan {
247            topology: decision.topology,
248            decision,
249        };
250        self.entries.insert(
251            key,
252            CudaMegakernelPlanCacheEntry {
253                plan,
254                last_seen: serial,
255            },
256        );
257        self.eviction_queue.push(Reverse((serial, key)));
258        self.update_latest_identity(key.identity(), serial, plan.topology);
259        Ok(plan)
260    }
261
262    /// Return a cached topology plan or select and cache one from the current
263    /// CUDA telemetry sample.
264    ///
265    /// This is the hot-path convenience API: callers provide stable graph,
266    /// analysis, device, and telemetry inputs, while the cache owns the
267    /// pressure bucketing needed to avoid stale sparse/dense decisions.
268    pub fn get_or_select_topology(
269        &mut self,
270        graph_layout_hash: u64,
271        analysis_kind: CudaMegakernelAnalysisKind,
272        device: CudaMegakernelDeviceKey,
273        sample: CudaMegakernelScheduleSample,
274        graph: MegakernelGraphShape,
275        memory: MegakernelMemoryBudget,
276        launch_overhead_ns: f64,
277        fusion_pressure: f64,
278    ) -> Result<CudaMegakernelCachedPlan, MegakernelMemoryError> {
279        let effective_fusion_pressure = if device.supports_grid_sync {
280            fusion_pressure
281        } else {
282            0.0
283        };
284        let key = CudaMegakernelPlanCacheKey::new(
285            graph_layout_hash,
286            analysis_kind,
287            device,
288            sample.frontier_density,
289            pressure_bps(memory.required_bytes, memory.budget_bytes),
290            sample.readback_bytes,
291            launch_pressure_bps(sample.dispatch_cost_ns, launch_overhead_ns),
292            effective_fusion_pressure,
293        );
294        let previous_topology =
295            self.latest_topology_for_identity(graph_layout_hash, analysis_kind, device);
296        self.get_or_insert_with(key, || {
297            if let Some(previous_topology) = previous_topology {
298                select_cuda_megakernel_topology_stable(
299                    sample,
300                    graph,
301                    memory,
302                    launch_overhead_ns,
303                    effective_fusion_pressure,
304                    previous_topology,
305                )
306            } else {
307                select_cuda_megakernel_topology(
308                    sample,
309                    graph,
310                    memory,
311                    launch_overhead_ns,
312                    effective_fusion_pressure,
313                )
314            }
315        })
316    }
317
318    /// Return a cache-backed, memory-validated CUDA megakernel execution plan.
319    ///
320    /// The cache key uses sparse-plan memory pressure because sparse is the
321    /// lower-bound resident footprint shared by every topology. A cache hit
322    /// reuses the prior topology decision, then this method validates the exact
323    /// current dense/fused/sparse byte budget before returning a launchable
324    /// plan. If the cached non-sparse topology no longer fits, the method
325    /// downgrades to sparse only after proving the sparse plan fits.
326    pub fn get_or_plan_execution(
327        &mut self,
328        graph_layout_hash: u64,
329        analysis_kind: CudaMegakernelAnalysisKind,
330        device: CudaMegakernelDeviceKey,
331        sample: CudaMegakernelScheduleSample,
332        graph: MegakernelGraphShape,
333        bytes_per_node: u64,
334        bytes_per_edge: u64,
335        frontier_bytes: u64,
336        scratch_bytes: u64,
337        output_bytes: u64,
338        budget_bytes: u64,
339        launch_overhead_ns: f64,
340        fusion_pressure: f64,
341    ) -> Result<MegakernelExecutionPlan, MegakernelMemoryError> {
342        let sparse_memory = plan_megakernel_memory_budget(
343            MegakernelExecutionTopology::SparseFrontier,
344            graph,
345            bytes_per_node,
346            bytes_per_edge,
347            frontier_bytes,
348            scratch_bytes,
349            output_bytes,
350            u64::MAX,
351        )?;
352        let cached = self.get_or_select_topology(
353            graph_layout_hash,
354            analysis_kind,
355            device,
356            sample,
357            graph,
358            MegakernelMemoryBudget {
359                required_bytes: sparse_memory.required_bytes,
360                budget_bytes,
361            },
362            launch_overhead_ns,
363            fusion_pressure,
364        )?;
365        match plan_megakernel_memory_budget(
366            cached.topology,
367            graph,
368            bytes_per_node,
369            bytes_per_edge,
370            frontier_bytes,
371            scratch_bytes,
372            output_bytes,
373            budget_bytes,
374        ) {
375            Ok(memory) => Ok(MegakernelExecutionPlan {
376                topology: cached.topology,
377                memory,
378                downgraded_to_sparse: false,
379            }),
380            Err(MegakernelMemoryError::OverBudget { .. })
381                if cached.topology != MegakernelExecutionTopology::SparseFrontier =>
382            {
383                let memory = plan_megakernel_memory_budget(
384                    MegakernelExecutionTopology::SparseFrontier,
385                    graph,
386                    bytes_per_node,
387                    bytes_per_edge,
388                    frontier_bytes,
389                    scratch_bytes,
390                    output_bytes,
391                    budget_bytes,
392                )?;
393                Ok(MegakernelExecutionPlan {
394                    topology: MegakernelExecutionTopology::SparseFrontier,
395                    memory,
396                    downgraded_to_sparse: true,
397                })
398            }
399            Err(error) => Err(error),
400        }
401    }
402
403    /// Return cache counters.
404    #[must_use]
405    pub fn stats(&self) -> CudaMegakernelPlanCacheStats {
406        CudaMegakernelPlanCacheStats {
407            hits: self.hits,
408            misses: self.misses,
409            evictions: self.evictions,
410            entries: self.entries.len(),
411        }
412    }
413
414    /// Drop every cached plan and preserve counters for observability.
415    pub fn clear(&mut self) {
416        self.entries.clear();
417        self.latest_by_identity.clear();
418        self.eviction_queue.clear();
419    }
420
421    fn latest_topology_for_identity(
422        &self,
423        graph_layout_hash: u64,
424        analysis_kind: CudaMegakernelAnalysisKind,
425        device: CudaMegakernelDeviceKey,
426    ) -> Option<MegakernelExecutionTopology> {
427        self.latest_by_identity
428            .get(&CudaMegakernelPlanIdentityKey {
429                graph_layout_hash,
430                analysis_kind,
431                device,
432            })
433            .map(|(_, topology)| *topology)
434    }
435
436    fn update_latest_identity(
437        &mut self,
438        identity: CudaMegakernelPlanIdentityKey,
439        serial: u64,
440        topology: MegakernelExecutionTopology,
441    ) {
442        match self.latest_by_identity.get(&identity) {
443            Some((latest_serial, _)) if *latest_serial > serial => {}
444            _ => {
445                self.latest_by_identity.insert(identity, (serial, topology));
446            }
447        }
448    }
449
450    fn recompute_latest_identity(&mut self, identity: CudaMegakernelPlanIdentityKey) {
451        let latest = self
452            .entries
453            .iter()
454            .filter(|(key, _)| key.identity() == identity)
455            .max_by_key(|(_, entry)| entry.last_seen)
456            .map(|(_, entry)| (entry.last_seen, entry.plan.topology));
457        if let Some(latest) = latest {
458            self.latest_by_identity.insert(identity, latest);
459        } else {
460            self.latest_by_identity.remove(&identity);
461        }
462    }
463
464    fn evict_until_below_limit(&mut self) -> Result<(), MegakernelMemoryError> {
465        while self.entries.len() >= self.max_entries {
466            let Some(Reverse((last_seen, lru_key))) = self.eviction_queue.pop() else {
467                break;
468            };
469            let Some(entry) = self.entries.get(&lru_key) else {
470                continue;
471            };
472            if entry.last_seen != last_seen {
473                continue;
474            }
475            let identity = lru_key.identity();
476            let evicted_topology = entry.plan.topology;
477            self.entries.remove(&lru_key);
478            if matches!(
479                self.latest_by_identity.get(&identity),
480                Some((latest_seen, latest_topology))
481                    if *latest_seen == last_seen && *latest_topology == evicted_topology
482            ) {
483                self.recompute_latest_identity(identity);
484            }
485            increment_plan_cache_counter(
486                &mut self.evictions,
487                "megakernel plan-cache eviction counter",
488            );
489        }
490        Ok(())
491    }
492
493    fn advance_serial(&mut self) -> Result<u64, MegakernelMemoryError> {
494        if let Some(next) = self.serial.checked_add(1) {
495            self.serial = next;
496            return Ok(next);
497        }
498        self.rebase_lru_serials()?;
499        self.serial =
500            self.serial
501                .checked_add(1)
502                .ok_or(MegakernelMemoryError::ByteCountOverflow {
503                    field: "megakernel plan-cache LRU serial after rebase",
504                })?;
505        Ok(self.serial)
506    }
507
508    fn rebase_lru_serials(&mut self) -> Result<(), MegakernelMemoryError> {
509        let mut ordered = Vec::new();
510        reserve_vec(
511            &mut ordered,
512            self.entries.len(),
513            "megakernel plan-cache LRU rebase scratch",
514        )
515        .map_err(|_| MegakernelMemoryError::ByteCountOverflow {
516            field: "megakernel plan-cache LRU rebase scratch",
517        })?;
518        for (key, entry) in &self.entries {
519            ordered.push((entry.last_seen, *key));
520        }
521        sort_unstable_by_key_if_needed(&mut ordered, |(last_seen, key)| (*last_seen, *key));
522        self.eviction_queue.clear();
523        self.latest_by_identity.clear();
524        let mut serial = 0_u64;
525        for (_, key) in ordered {
526            serial = serial
527                .checked_add(1)
528                .ok_or(MegakernelMemoryError::ByteCountOverflow {
529                    field: "megakernel plan-cache LRU rebase serial",
530                })?;
531            let topology = if let Some(entry) = self.entries.get_mut(&key) {
532                entry.last_seen = serial;
533                Some(entry.plan.topology)
534            } else {
535                None
536            };
537            if let Some(topology) = topology {
538                self.eviction_queue.push(Reverse((serial, key)));
539                self.update_latest_identity(key.identity(), serial, topology);
540            }
541        }
542        self.serial = serial;
543        Ok(())
544    }
545}
546
547fn density_bucket(frontier_density: f64) -> u16 {
548    if !frontier_density.is_finite() {
549        return 0;
550    }
551    let clamped = frontier_density.clamp(0.0, 1.0);
552    rounded_f64_to_u16_bucket(
553        clamped * f64::from(DENSITY_BUCKETS - 1),
554        "frontier-density bucket",
555    )
556}
557
558fn pressure_bucket(memory_pressure_bps: u32) -> u32 {
559    memory_pressure_bps / PRESSURE_BUCKET_BPS
560}
561
562fn pressure_bps(numerator: u64, denominator: u64) -> u32 {
563    crate::numeric::CUDA_NUMERIC.ratio_basis_points_u64(
564        numerator,
565        denominator,
566        if numerator == 0 { 0 } else { u32::MAX },
567        "megakernel pressure",
568    )
569}
570
571fn launch_pressure_bps(dispatch_cost_ns: f64, launch_overhead_ns: f64) -> u32 {
572    crate::numeric::CUDA_NUMERIC.finite_f64_ratio_basis_points_trunc(
573        launch_overhead_ns,
574        dispatch_cost_ns,
575        u32::MAX,
576        0,
577        "launch-pressure basis-points",
578    )
579}
580
581fn readback_bucket(readback_bytes: u64) -> u16 {
582    if readback_bytes == 0 {
583        return 0;
584    }
585    let shifted = readback_bytes >> READBACK_BUCKET_SHIFT;
586    let bucket = u64::BITS - shifted.leading_zeros();
587    bucket.min(u32::from(u16::MAX)) as u16
588}
589
590fn fusion_bucket(fusion_pressure: f64) -> u32 {
591    pressure_bucket(
592        crate::numeric::CUDA_NUMERIC.finite_f64_unit_basis_points_trunc(
593            fusion_pressure,
594            0,
595            "fusion-pressure basis-points",
596        ),
597    )
598}
599
600fn rounded_f64_to_u16_bucket(value: f64, label: &'static str) -> u16 {
601    let rounded = value.round();
602    if !rounded.is_finite() || rounded < 0.0 || rounded > f64::from(u16::MAX) {
603        tracing::error!(
604            "CUDA megakernel {label} value {rounded} cannot fit u16. Fix: reduce bucket resolution or shard cache domains."
605        );
606        return if rounded.is_sign_negative() {
607            0
608        } else {
609            u16::MAX
610        };
611    }
612    rounded as u16
613}
614
615#[cfg(test)]
616mod tests {
617    use super::{
618        CudaMegakernelAnalysisKind, CudaMegakernelDeviceKey, CudaMegakernelPlanCache,
619        CudaMegakernelPlanCacheKey,
620    };
621    use crate::megakernel_scheduler::CudaMegakernelScheduleSample;
622    use crate::synthetic_device_caps::synthetic_sm120_envelope_default;
623    use vyre_driver::megakernel_execution::{
624        MegakernelExecutionTopology, MegakernelGraphShape, MegakernelTopologyDecision,
625    };
626
627    fn device() -> CudaMegakernelDeviceKey {
628        CudaMegakernelDeviceKey {
629            sm_major: 12,
630            sm_minor: 0,
631            warp_size: 32,
632            supports_grid_sync: true,
633            supports_tensor_cores: true,
634            max_workgroup_size: 1024,
635        }
636    }
637
638    fn key(
639        graph_layout_hash: u64,
640        analysis_kind: CudaMegakernelAnalysisKind,
641        frontier_density: f64,
642        memory_pressure_bps: u32,
643    ) -> CudaMegakernelPlanCacheKey {
644        CudaMegakernelPlanCacheKey::new(
645            graph_layout_hash,
646            analysis_kind,
647            device(),
648            frontier_density,
649            memory_pressure_bps,
650            0,
651            0,
652            0.0,
653        )
654    }
655
656    fn decision(topology: MegakernelExecutionTopology) -> MegakernelTopologyDecision {
657        MegakernelTopologyDecision {
658            topology,
659            memory_pressure_bps: 1_000,
660            average_degree_bps: 20_000,
661            launch_pressure_bps: 2_000,
662        }
663    }
664
665    #[test]
666    fn cache_reuses_plan_for_same_graph_analysis_device_and_pressure_bucket() {
667        let mut cache = CudaMegakernelPlanCache::new();
668        let key = key(42, CudaMegakernelAnalysisKind::Ifds, 0.52, 2_400);
669        let first = cache
670            .get_or_insert_with(key, || decision(MegakernelExecutionTopology::FusedWave))
671            .expect("Fix: CUDA megakernel plan-cache insert should fit telemetry counters.");
672        let second = cache
673            .get_or_insert_with(key, || {
674                decision(MegakernelExecutionTopology::SparseFrontier)
675            })
676            .expect("Fix: CUDA megakernel plan-cache hit should fit telemetry counters.");
677
678        assert_eq!(first, second);
679        assert_eq!(second.topology, MegakernelExecutionTopology::FusedWave);
680        let stats = cache.stats();
681        assert_eq!(stats.hits, 1);
682        assert_eq!(stats.misses, 1);
683        assert_eq!(stats.entries, 1);
684    }
685
686    #[test]
687    fn device_key_is_derived_from_cuda_caps() {
688        assert_eq!(
689            CudaMegakernelDeviceKey::from(&synthetic_sm120_envelope_default()),
690            device()
691        );
692    }
693
694    #[test]
695    fn cache_separates_analysis_family_density_and_device_features() {
696        let ifds = key(42, CudaMegakernelAnalysisKind::Ifds, 0.01, 1_000);
697        let liveness = key(42, CudaMegakernelAnalysisKind::Liveness, 0.01, 1_000);
698        let dense = key(42, CudaMegakernelAnalysisKind::Ifds, 0.95, 1_000);
699        let mut other_device = device();
700        other_device.sm_minor = 1;
701        let device_changed = CudaMegakernelPlanCacheKey::new(
702            42,
703            CudaMegakernelAnalysisKind::Ifds,
704            other_device,
705            0.01,
706            1_000,
707            0,
708            0,
709            0.0,
710        );
711
712        assert_ne!(ifds, liveness);
713        assert_ne!(ifds, dense);
714        assert_ne!(ifds, device_changed);
715    }
716
717    #[test]
718    fn bounded_cache_evicts_lru_entry() {
719        let mut cache = CudaMegakernelPlanCache::with_max_entries(2);
720        let first = key(1, CudaMegakernelAnalysisKind::Dataflow, 0.1, 1_000);
721        let second = key(2, CudaMegakernelAnalysisKind::Dataflow, 0.1, 1_000);
722        let third = key(3, CudaMegakernelAnalysisKind::Dataflow, 0.1, 1_000);
723
724        cache
725            .get_or_insert_with(first, || {
726                decision(MegakernelExecutionTopology::SparseFrontier)
727            })
728            .expect("Fix: CUDA megakernel plan-cache insert should fit telemetry counters.");
729        cache
730            .get_or_insert_with(second, || {
731                decision(MegakernelExecutionTopology::HybridFrontier)
732            })
733            .expect("Fix: CUDA megakernel plan-cache insert should fit telemetry counters.");
734        cache
735            .get_or_insert_with(first, || {
736                decision(MegakernelExecutionTopology::DenseFrontier)
737            })
738            .expect("Fix: CUDA megakernel plan-cache hit should fit telemetry counters.");
739        cache
740            .get_or_insert_with(third, || decision(MegakernelExecutionTopology::FusedWave))
741            .expect("Fix: CUDA megakernel plan-cache eviction should fit telemetry counters.");
742
743        let stats = cache.stats();
744        assert_eq!(stats.hits, 1);
745        assert_eq!(stats.misses, 3);
746        assert_eq!(stats.evictions, 1);
747        assert_eq!(stats.entries, 2);
748        let reloaded_second = cache
749            .get_or_insert_with(second, || {
750                decision(MegakernelExecutionTopology::DenseFrontier)
751            })
752            .expect("Fix: CUDA megakernel plan-cache reload should fit telemetry counters.");
753        assert_eq!(
754            reloaded_second.topology,
755            MegakernelExecutionTopology::DenseFrontier
756        );
757    }
758
759    #[test]
760    fn cache_selects_topology_and_reuses_pressure_bucket_plan() {
761        let mut cache = CudaMegakernelPlanCache::new();
762        let sample = crate::megakernel_scheduler::CudaMegakernelScheduleSample {
763            dispatch_cost_ns: 1_000.0,
764            frontier_density: 0.90,
765            readback_bytes: 1 << 20,
766        };
767        let graph = vyre_driver::megakernel_execution::MegakernelGraphShape {
768            node_count: 1_000,
769            edge_count: 4_000,
770        };
771        let memory = vyre_driver::megakernel_execution::MegakernelMemoryBudget {
772            required_bytes: 1_024,
773            budget_bytes: 16_384,
774        };
775        let first = cache
776            .get_or_select_topology(
777                99,
778                CudaMegakernelAnalysisKind::Dataflow,
779                device(),
780                sample,
781                graph,
782                memory,
783                250.0,
784                0.95,
785            )
786            .expect("Fix: CUDA megakernel topology selection should fit telemetry counters.");
787        let second = cache
788            .get_or_select_topology(
789                99,
790                CudaMegakernelAnalysisKind::Dataflow,
791                device(),
792                crate::megakernel_scheduler::CudaMegakernelScheduleSample {
793                    frontier_density: 0.91,
794                    ..sample
795                },
796                graph,
797                vyre_driver::megakernel_execution::MegakernelMemoryBudget {
798                    required_bytes: 1_100,
799                    budget_bytes: 16_384,
800                },
801                250.0,
802                0.95,
803            )
804            .expect("Fix: CUDA megakernel topology cache hit should fit telemetry counters.");
805
806        assert_eq!(first, second);
807        assert_eq!(first.topology, MegakernelExecutionTopology::FusedWave);
808        assert_eq!(cache.stats().hits, 1);
809        assert_eq!(cache.stats().misses, 1);
810    }
811
812    #[test]
813    fn cache_stabilizes_topology_across_adjacent_pressure_buckets() {
814        let mut cache = CudaMegakernelPlanCache::new();
815        let graph = vyre_driver::megakernel_execution::MegakernelGraphShape {
816            node_count: 1_000,
817            edge_count: 4_000,
818        };
819        let memory = vyre_driver::megakernel_execution::MegakernelMemoryBudget {
820            required_bytes: 1_024,
821            budget_bytes: 16_384,
822        };
823        let dense = cache
824            .get_or_select_topology(
825                99,
826                CudaMegakernelAnalysisKind::Dataflow,
827                device(),
828                crate::megakernel_scheduler::CudaMegakernelScheduleSample {
829                    dispatch_cost_ns: 1_000.0,
830                    frontier_density: 0.70,
831                    readback_bytes: 512,
832                },
833                graph,
834                memory,
835                100.0,
836                0.0,
837            )
838            .expect("Fix: CUDA megakernel topology selection should fit telemetry counters.");
839        let near_dense = cache
840            .get_or_select_topology(
841                99,
842                CudaMegakernelAnalysisKind::Dataflow,
843                device(),
844                crate::megakernel_scheduler::CudaMegakernelScheduleSample {
845                    dispatch_cost_ns: 1_000.0,
846                    frontier_density: 0.68,
847                    readback_bytes: 512,
848                },
849                graph,
850                memory,
851                100.0,
852                0.0,
853            )
854            .expect("Fix: CUDA megakernel topology stabilization should fit telemetry counters.");
855
856        assert_eq!(dense.topology, MegakernelExecutionTopology::DenseFrontier);
857        assert_eq!(
858            near_dense.topology,
859            MegakernelExecutionTopology::DenseFrontier
860        );
861        assert_eq!(cache.stats().hits, 0);
862        assert_eq!(cache.stats().misses, 2);
863    }
864
865    #[test]
866    fn cache_reselects_when_memory_pressure_bucket_changes() {
867        let mut cache = CudaMegakernelPlanCache::new();
868        let sample = crate::megakernel_scheduler::CudaMegakernelScheduleSample {
869            dispatch_cost_ns: 1_000.0,
870            frontier_density: 0.90,
871            readback_bytes: 1 << 20,
872        };
873        let graph = vyre_driver::megakernel_execution::MegakernelGraphShape {
874            node_count: 1_000,
875            edge_count: 4_000,
876        };
877        let low_pressure = cache
878            .get_or_select_topology(
879                99,
880                CudaMegakernelAnalysisKind::Dataflow,
881                device(),
882                sample,
883                graph,
884                vyre_driver::megakernel_execution::MegakernelMemoryBudget {
885                    required_bytes: 1_024,
886                    budget_bytes: 16_384,
887                },
888                250.0,
889                0.95,
890            )
891            .expect("Fix: CUDA megakernel topology selection should fit telemetry counters.");
892        let red_zone = cache
893            .get_or_select_topology(
894                99,
895                CudaMegakernelAnalysisKind::Dataflow,
896                device(),
897                sample,
898                graph,
899                vyre_driver::megakernel_execution::MegakernelMemoryBudget {
900                    required_bytes: 15_500,
901                    budget_bytes: 16_384,
902                },
903                250.0,
904                0.95,
905            )
906            .expect("Fix: CUDA megakernel topology reselection should fit telemetry counters.");
907
908        assert_eq!(
909            low_pressure.topology,
910            MegakernelExecutionTopology::FusedWave
911        );
912        assert_eq!(
913            red_zone.topology,
914            MegakernelExecutionTopology::SparseFrontier
915        );
916        assert_eq!(cache.stats().hits, 0);
917        assert_eq!(cache.stats().misses, 2);
918    }
919
920    #[test]
921    fn cache_pressure_bucket_uses_exact_u128_math() {
922        let low = CudaMegakernelPlanCacheKey::new(
923            1,
924            CudaMegakernelAnalysisKind::Dataflow,
925            device(),
926            0.5,
927            super::pressure_bps(1_u64 << 62, 1_u64 << 63),
928            0,
929            0,
930            0.0,
931        );
932        let high = CudaMegakernelPlanCacheKey::new(
933            1,
934            CudaMegakernelAnalysisKind::Dataflow,
935            device(),
936            0.5,
937            super::pressure_bps(1_u64 << 63, 1_u64 << 63),
938            0,
939            0,
940            0.0,
941        );
942
943        assert_eq!(low.memory_pressure_bucket, 5);
944        assert_eq!(high.memory_pressure_bucket, 10);
945    }
946
947    #[test]
948    fn cache_reselects_when_readback_launch_or_fusion_pressure_changes() {
949        let mut cache = CudaMegakernelPlanCache::new();
950        let graph = MegakernelGraphShape {
951            node_count: 1_000,
952            edge_count: 4_000,
953        };
954        let memory = vyre_driver::megakernel_execution::MegakernelMemoryBudget {
955            required_bytes: 1_024,
956            budget_bytes: 16_384,
957        };
958        let low_pressure = cache
959            .get_or_select_topology(
960                99,
961                CudaMegakernelAnalysisKind::Dataflow,
962                device(),
963                CudaMegakernelScheduleSample {
964                    dispatch_cost_ns: 1_000.0,
965                    frontier_density: 0.50,
966                    readback_bytes: 0,
967                },
968                graph,
969                memory,
970                250.0,
971                0.95,
972            )
973            .expect("Fix: CUDA megakernel topology selection should fit telemetry counters.");
974        let high_pressure = cache
975            .get_or_select_topology(
976                99,
977                CudaMegakernelAnalysisKind::Dataflow,
978                device(),
979                CudaMegakernelScheduleSample {
980                    dispatch_cost_ns: 1_000.0,
981                    frontier_density: 0.50,
982                    readback_bytes: 1 << 20,
983                },
984                graph,
985                memory,
986                250.0,
987                0.95,
988            )
989            .expect("Fix: CUDA megakernel topology pressure split should fit telemetry counters.");
990
991        assert_ne!(
992            low_pressure.topology,
993            MegakernelExecutionTopology::FusedWave
994        );
995        assert_eq!(
996            high_pressure.topology,
997            MegakernelExecutionTopology::FusedWave
998        );
999        assert_eq!(cache.stats().hits, 0);
1000        assert_eq!(cache.stats().misses, 2);
1001    }
1002
1003    #[test]
1004    fn cache_never_selects_fused_wave_without_grid_sync_support() {
1005        let mut cache = CudaMegakernelPlanCache::new();
1006        let mut no_grid_sync = device();
1007        no_grid_sync.supports_grid_sync = false;
1008
1009        let plan = cache
1010            .get_or_select_topology(
1011                99,
1012                CudaMegakernelAnalysisKind::Dataflow,
1013                no_grid_sync,
1014                CudaMegakernelScheduleSample {
1015                    dispatch_cost_ns: 1_000.0,
1016                    frontier_density: 0.50,
1017                    readback_bytes: 1 << 20,
1018                },
1019                MegakernelGraphShape {
1020                    node_count: 1_000,
1021                    edge_count: 4_000,
1022                },
1023                vyre_driver::megakernel_execution::MegakernelMemoryBudget {
1024                    required_bytes: 1_024,
1025                    budget_bytes: 16_384,
1026                },
1027                250.0,
1028                0.95,
1029            )
1030            .expect("Fix: CUDA megakernel topology selection should fit telemetry counters.");
1031
1032        assert_ne!(
1033            plan.topology,
1034            MegakernelExecutionTopology::FusedWave,
1035            "Fix: CUDA megakernel planner must not select cooperative fused-wave topology when the device key says grid sync is unavailable."
1036        );
1037    }
1038
1039    #[test]
1040    fn cached_execution_plan_reuses_topology_bucket_and_validates_memory() {
1041        let mut cache = CudaMegakernelPlanCache::new();
1042        let sample = CudaMegakernelScheduleSample {
1043            dispatch_cost_ns: 1_000.0,
1044            frontier_density: 0.90,
1045            readback_bytes: 1 << 20,
1046        };
1047        let graph = MegakernelGraphShape {
1048            node_count: 1_000,
1049            edge_count: 4_000,
1050        };
1051        let first = cache
1052            .get_or_plan_execution(
1053                99,
1054                CudaMegakernelAnalysisKind::Dataflow,
1055                device(),
1056                sample,
1057                graph,
1058                16,
1059                8,
1060                4_096,
1061                2_048,
1062                512,
1063                128 * 1024,
1064                250.0,
1065                0.95,
1066            )
1067            .expect("Fix: cache-backed fused CUDA execution plan should fit the explicit budget.");
1068        let second = cache
1069            .get_or_plan_execution(
1070                99,
1071                CudaMegakernelAnalysisKind::Dataflow,
1072                device(),
1073                CudaMegakernelScheduleSample {
1074                    frontier_density: 0.91,
1075                    ..sample
1076                },
1077                graph,
1078                16,
1079                8,
1080                4_096,
1081                2_048,
1082                512,
1083                128 * 1024,
1084                250.0,
1085                0.95,
1086            )
1087            .expect("Fix: equivalent CUDA execution pressure bucket should reuse the cached topology and still validate memory.");
1088
1089        assert_eq!(first.topology, MegakernelExecutionTopology::FusedWave);
1090        assert_eq!(second.topology, MegakernelExecutionTopology::FusedWave);
1091        assert_eq!(second.memory.scratch_bytes, 8_192);
1092        assert!(!second.downgraded_to_sparse);
1093        assert_eq!(cache.stats().hits, 1);
1094        assert_eq!(cache.stats().misses, 1);
1095    }
1096
1097    #[test]
1098    fn cached_execution_plan_downgrades_non_sparse_topology_when_exact_budget_fails() {
1099        let mut cache = CudaMegakernelPlanCache::new();
1100        let plan = cache
1101            .get_or_plan_execution(
1102                99,
1103                CudaMegakernelAnalysisKind::Dataflow,
1104                device(),
1105                CudaMegakernelScheduleSample {
1106                    dispatch_cost_ns: 1_000.0,
1107                    frontier_density: 0.50,
1108                    readback_bytes: 1 << 20,
1109                },
1110                MegakernelGraphShape {
1111                    node_count: 1_000,
1112                    edge_count: 4_000,
1113                },
1114                16,
1115                8,
1116                4_096,
1117                10_000,
1118                512,
1119                80_000,
1120                250.0,
1121                0.90,
1122            )
1123            .expect("Fix: sparse CUDA downgrade must fit after cached fused topology exceeds exact budget.");
1124
1125        assert_eq!(plan.topology, MegakernelExecutionTopology::SparseFrontier);
1126        assert!(plan.downgraded_to_sparse);
1127        assert_eq!(plan.memory.scratch_bytes, 10_000);
1128        assert_eq!(cache.stats().misses, 1);
1129        assert_eq!(cache.stats().entries, 1);
1130    }
1131
1132    #[test]
1133    fn cache_rebases_lru_serial_instead_of_failing_dispatch() {
1134        let mut cache = CudaMegakernelPlanCache::with_max_entries(2);
1135        let first = key(1, CudaMegakernelAnalysisKind::Ifds, 0.10, 1_000);
1136        let second = key(2, CudaMegakernelAnalysisKind::Ifds, 0.20, 1_000);
1137        cache
1138            .get_or_insert_with(first, || {
1139                decision(MegakernelExecutionTopology::SparseFrontier)
1140            })
1141            .expect("Fix: first plan insert should fit");
1142        cache
1143            .get_or_insert_with(second, || {
1144                decision(MegakernelExecutionTopology::DenseFrontier)
1145            })
1146            .expect("Fix: second plan insert should fit");
1147        cache.serial = u64::MAX;
1148
1149        cache
1150            .get_or_insert_with(first, || decision(MegakernelExecutionTopology::FusedWave))
1151            .expect(
1152                "Fix: LRU serial exhaustion must rebase instead of failing the CUDA dispatch path",
1153            );
1154
1155        let first_seen = cache
1156            .entries
1157            .get(&first)
1158            .expect("Fix: first entry must remain")
1159            .last_seen;
1160        let second_seen = cache
1161            .entries
1162            .get(&second)
1163            .expect("Fix: second entry must remain")
1164            .last_seen;
1165        assert!(first_seen > second_seen);
1166        assert_eq!(cache.stats().hits, 1);
1167    }
1168
1169    #[test]
1170    fn cache_counters_pin_instead_of_failing_dispatch() {
1171        let mut cache = CudaMegakernelPlanCache::new();
1172        let key = key(3, CudaMegakernelAnalysisKind::Ifds, 0.10, 1_000);
1173        cache
1174            .get_or_insert_with(key, || {
1175                decision(MegakernelExecutionTopology::SparseFrontier)
1176            })
1177            .expect("Fix: plan insert should fit");
1178        cache.hits = u64::MAX;
1179
1180        cache
1181            .get_or_insert_with(key, || decision(MegakernelExecutionTopology::DenseFrontier))
1182            .expect("Fix: counter exhaustion must not fail the CUDA dispatch path");
1183
1184        assert_eq!(cache.stats().hits, u64::MAX);
1185    }
1186}
vyre_driver_cuda/megakernel_plan_cache.rs

vyre_driver_cuda/
megakernel_plan_cache.rs
