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 assigns_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 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_writes.entry(name.to_string()).or_insert(0) += 1;
70 }
71
72 _ => {}
73 }
74 }
75
76 let mut function_locals = Vec::new();
77 if let Ok(module) = vyre_emit_naga::emit(desc) {
78 for ep in &module.entry_points {
79 for (_, local) in ep.function.local_variables.iter() {
80 if let Some(name) = &local.name {
81 if name.starts_with("vyre_named_carry_") {
82 function_locals.push(name.clone());
83 }
84 }
85 }
86 }
87 }
88
89 CarrierSummary {
90 assigns_observed: desc.ops_iter().count(),
91 carrier_reads,
92 carrier_writes,
93 carrier_finals,
94 function_locals,
95 }
96}