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}