Skip to main content

vyre_debug/
descriptor_dump.rs

1use std::collections::BTreeMap;
2use std::fmt::Write;
3use vyre_lower::{BindingVisibility, KernelBody, KernelDescriptor};
4
5pub struct DescriptorDumpOptions {
6    pub show_literals: bool,
7    pub show_result_ids: bool,
8    pub max_ops_per_body: usize,
9}
10
11impl Default for DescriptorDumpOptions {
12    fn default() -> Self {
13        Self {
14            show_literals: true,
15            show_result_ids: true,
16            max_ops_per_body: usize::MAX,
17        }
18    }
19}
20
21#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
22pub struct DescriptorDump {
23    pub text: String,
24    #[serde(
25        serialize_with = "crate::path_map_serde::serialize_usize",
26        deserialize_with = "crate::path_map_serde::deserialize_usize"
27    )]
28    pub op_counts_by_path: BTreeMap<Vec<usize>, usize>,
29}
30
31pub fn dump_descriptor(desc: &KernelDescriptor, options: &DescriptorDumpOptions) -> DescriptorDump {
32    let mut out = String::new();
33    let mut counts = BTreeMap::new();
34
35    let id_prefix = if desc.id.len() > 8 {
36        &desc.id[0..8]
37    } else {
38        &desc.id
39    };
40    let _ = writeln!(out, "KernelDescriptor id={}", id_prefix);
41
42    let _ = writeln!(out, "bindings:");
43    for slot in &desc.bindings.slots {
44        let count_str = match slot.element_count {
45            Some(c) => format!("Some({})", c),
46            None => "None".to_string(),
47        };
48        let access_str = match slot.visibility {
49            BindingVisibility::ReadOnly => "ReadOnly",
50            BindingVisibility::ReadWrite => "ReadWrite",
51            BindingVisibility::WriteOnly => "WriteOnly",
52        };
53        // Just format it properly
54        let _ = writeln!(
55            out,
56            "  slot={} name={} count={} access={} mc={:?}",
57            slot.slot, slot.name, count_str, access_str, slot.memory_class
58        );
59    }
60
61    let _ = writeln!(
62        out,
63        "dispatch: workgroup_size={:?}",
64        desc.dispatch.workgroup_size
65    );
66
67    fn walk_body(
68        body: &KernelBody,
69        path: Vec<usize>,
70        indent: usize,
71        out: &mut String,
72        counts: &mut BTreeMap<Vec<usize>, usize>,
73        options: &DescriptorDumpOptions,
74    ) {
75        counts.insert(path.clone(), body.ops.len());
76
77        let path_str = if path.is_empty() {
78            "[]".to_string()
79        } else {
80            let s: Vec<_> = path.iter().map(|n| n.to_string()).collect();
81            format!("[{}]", s.join(","))
82        };
83
84        let indent_str = " ".repeat(indent);
85        let _ = writeln!(out, "{}body{}:", indent_str, path_str);
86
87        let mut ops_shown = 0;
88        for (i, op) in body.ops.iter().enumerate() {
89            if ops_shown >= options.max_ops_per_body {
90                let remaining = body.ops.len() - ops_shown;
91                let _ = writeln!(out, "{}  ... <{} more ops>", indent_str, remaining);
92                break;
93            }
94
95            let result_str = if options.show_result_ids {
96                match op.result {
97                    Some(r) => format!(" result=Some({})", r),
98                    None => " result=None".to_string(),
99                }
100            } else {
101                "".to_string()
102            };
103
104            let _ = writeln!(
105                out,
106                "{}  [{}] {:?} ops={:?}{}",
107                indent_str, i, op.kind, op.operands, result_str
108            );
109
110            // If the op kind has child bodies, we need to iterate over them.
111            // Wait, child_bodies is a flat array in KernelBody. The op just references them by index?
112            // Let's assume ops reference child body indices in `op.operands` or something?
113            // No, the dump says: "body[0,0,12,0]" for structured_if_then ops=[10, 0] -> actually vyre_lower has child bodies. Let me check KernelOp / KernelBody.
114            ops_shown += 1;
115        }
116
117        for (i, child) in body.child_bodies.iter().enumerate() {
118            let mut child_path = path.clone();
119            child_path.push(i);
120            walk_body(child, child_path, indent + 2, out, counts, options);
121        }
122    }
123
124    walk_body(&desc.body, vec![], 0, &mut out, &mut counts, options);
125
126    if options.show_literals {
127        let _ = writeln!(out, "literals:");
128        fn walk_literals(body: &KernelBody, path: Vec<usize>, out: &mut String) {
129            let path_str = if path.is_empty() {
130                "[]".to_string()
131            } else {
132                let s: Vec<_> = path.iter().map(|n| n.to_string()).collect();
133                format!("[{}]", s.join(","))
134            };
135
136            // Just formatting literals
137            let _ = writeln!(out, "  body{}: {:?}", path_str, body.literals);
138
139            for (i, child) in body.child_bodies.iter().enumerate() {
140                let mut child_path = path.clone();
141                child_path.push(i);
142                walk_literals(child, child_path, out);
143            }
144        }
145        walk_literals(&desc.body, vec![], &mut out);
146    }
147
148    // Remove trailing newline if any, or just keep it? Dump usually expects exact matches.
149    DescriptorDump {
150        text: out,
151        op_counts_by_path: counts,
152    }
153}