vyre-megakernel 0.7.2

Backend-neutral compiler for canonical Vyre megakernel artifacts
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

use vyre_foundation::ir::ProgramGraph;

use crate::{
    candidate::CandidatePlan,
    facts::{DataflowEdge, PlanningFacts},
    legality::{analyze_fusion_pair, FusionDecision, FusionRejectionReason},
    DependencyEdge, FusionGroupId, SearchBudget, SearchWork,
};

#[derive(Debug)]
pub(crate) struct RejectedEdge {
    pub(crate) edge: DataflowEdge,
    pub(crate) reason: FusionRejectionReason,
}

#[derive(Debug)]
pub(crate) struct SearchResult {
    pub(crate) candidates: Vec<CandidatePlan>,
    pub(crate) rejected: Vec<RejectedEdge>,
    pub(crate) work: SearchWork,
}

pub(crate) fn explore(
    graph: &ProgramGraph,
    facts: &PlanningFacts,
    dependencies: &[DependencyEdge],
    budget: SearchBudget,
) -> SearchResult {
    let mut candidates = vec![CandidatePlan::baseline(graph.nodes().len())];
    let mut rejected = Vec::new();
    let mut legal_edges = Vec::new();
    let mut cpu_work = 0_u64;

    for edge in &facts.dataflow {
        if !can_spend(cpu_work, budget) {
            break;
        }
        cpu_work = cpu_work.saturating_add(1);
        match analyze_fusion_pair(graph, edge.from, edge.to, edge.value) {
            FusionDecision::Legal => {
                let candidate = CandidatePlan::from_edges(graph.nodes().len(), &[*edge]);
                if candidate_is_acyclic(&candidate, dependencies) {
                    legal_edges.push(*edge);
                    if candidates.len() < budget.max_candidates as usize {
                        candidates.push(candidate);
                    }
                } else {
                    rejected.push(RejectedEdge {
                        edge: *edge,
                        reason: FusionRejectionReason::DependencyCycle,
                    });
                }
            }
            FusionDecision::Rejected(reason) => rejected.push(RejectedEdge {
                edge: *edge,
                reason,
            }),
        }
    }

    if legal_edges.len() > 1
        && candidates.len() < budget.max_candidates as usize
        && can_spend(cpu_work, budget)
    {
        cpu_work = cpu_work.saturating_add(1);
        let mut accepted = Vec::new();
        for edge in legal_edges {
            let mut proposed = accepted.clone();
            proposed.push(edge);
            let candidate = CandidatePlan::from_edges(graph.nodes().len(), &proposed);
            if candidate_is_acyclic(&candidate, dependencies) {
                accepted = proposed;
            } else {
                rejected.push(RejectedEdge {
                    edge,
                    reason: FusionRejectionReason::DependencyCycle,
                });
            }
        }
        candidates.push(CandidatePlan::from_edges(graph.nodes().len(), &accepted));
    }
    candidates.sort_by(|left, right| left.node_groups.cmp(&right.node_groups));
    candidates.dedup_by(|left, right| left.node_groups == right.node_groups);

    SearchResult {
        work: SearchWork {
            candidates_explored: u32::try_from(candidates.len()).unwrap_or(u32::MAX),
            cpu_work,
            target_compilations: 0,
            measurements: 0,
            elapsed_ns: cpu_work.min(budget.max_elapsed_ns),
        },
        candidates,
        rejected,
    }
}

fn candidate_is_acyclic(candidate: &CandidatePlan, dependencies: &[DependencyEdge]) -> bool {
    let groups = candidate
        .node_groups
        .iter()
        .copied()
        .map(FusionGroupId)
        .collect::<Vec<_>>();
    crate::group_stages(candidate.group_count(), dependencies, &groups).is_ok()
}

fn can_spend(cpu_work: u64, budget: SearchBudget) -> bool {
    cpu_work < budget.max_cpu_work && cpu_work < budget.max_elapsed_ns
}