Skip to main content

vyre_debug/
carriers.rs

1use serde::{Deserialize, Serialize};
2use std::collections::BTreeMap;
3use vyre_foundation::ir::Program;
4use vyre_lower::{KernelDescriptor, KernelOpKind};
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct UncarrieredAssign {
8    pub name: String,
9    pub loop_path: Vec<String>,
10    pub has_carrier_op: bool,
11    pub has_final_op: bool,
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct CarrierSummary {
16    pub total_ops_observed: usize,
17    pub carrier_reads: BTreeMap<String, usize>,
18    pub carrier_writes: BTreeMap<String, usize>,
19    pub carrier_finals: BTreeMap<String, usize>,
20    pub function_locals: Vec<String>,
21}
22
23pub fn find_uncarriered_assigns(
24    program: &Program,
25    desc: &KernelDescriptor,
26) -> Vec<UncarrieredAssign> {
27    let mut assigns = Vec::new();
28
29    crate::source_walker::walk_source_assigns(program, |name, loop_path| {
30        let mut has_carrier_op = false;
31        let mut has_final_op = false;
32
33        for op in desc.ops_iter() {
34            match &op.kind {
35                KernelOpKind::LoopCarrier { name: n } if n.as_ref() == name => {
36                    has_carrier_op = true
37                }
38                KernelOpKind::LoopCarrierEnd { name: n } if n.as_ref() == name => {
39                    has_final_op = true
40                }
41                _ => {}
42            }
43        }
44
45        if !has_carrier_op || !has_final_op {
46            assigns.push(UncarrieredAssign {
47                name: name.to_string(),
48                loop_path: loop_path.clone(),
49                has_carrier_op,
50                has_final_op,
51            });
52        }
53    });
54
55    assigns
56}
57
58pub fn carrier_summary(desc: &KernelDescriptor) -> CarrierSummary {
59    let mut carrier_reads = BTreeMap::new();
60    let mut carrier_writes = BTreeMap::new();
61    let mut carrier_finals = BTreeMap::new();
62
63    for op in desc.ops_iter() {
64        match &op.kind {
65            KernelOpKind::LoopCarrier { name } => {
66                *carrier_reads.entry(name.to_string()).or_insert(0) += 1;
67            }
68            KernelOpKind::LoopCarrierEnd { name } => {
69                *carrier_finals.entry(name.to_string()).or_insert(0) += 1;
70            }
71            KernelOpKind::LoopCarrierInit { name } => {
72                *carrier_writes.entry(name.to_string()).or_insert(0) += 1;
73            }
74            _ => {}
75        }
76    }
77
78    let mut function_locals = Vec::new();
79    if let Ok(module) = vyre_emit_naga::emit(desc) {
80        for ep in &module.entry_points {
81            for (_, local) in ep.function.local_variables.iter() {
82                if let Some(name) = &local.name {
83                    if name.starts_with("vyre_named_carry_") {
84                        function_locals.push(name.clone());
85                    }
86                }
87            }
88        }
89    }
90
91    CarrierSummary {
92        total_ops_observed: desc.ops_iter().count(),
93        carrier_reads,
94        carrier_writes,
95        carrier_finals,
96        function_locals,
97    }
98}