cubecl_opt/
debug.rs

1use core::fmt::Display;
2
3use alloc::{
4    format,
5    rc::Rc,
6    string::{String, ToString},
7    vec::Vec,
8};
9use petgraph::{
10    dot::{Config, Dot},
11    prelude::StableDiGraph,
12    visit::EdgeRef,
13};
14
15use crate::{
16    BasicBlock, ControlFlow, Function, Optimizer,
17    analyses::{
18        liveness::{
19            Liveness,
20            shared::{SharedLiveness, SmemAllocation},
21        },
22        uniformity::Uniformity,
23    },
24    gvn::{BlockSets, Expression, GlobalValues, Instruction, ValueTable},
25};
26
27const DEBUG_GVN: bool = option_env!("CUBECL_DEBUG_GVN").is_some();
28
29impl Display for Optimizer {
30    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
31        write!(f, "main: {{\n{}\n}}", self.main)?;
32        for (id, extra_func) in self.global_state.extra_functions.iter() {
33            write!(
34                f,
35                "\n\nfunc_{id}[{}]({}): {{\n{}\n}}",
36                extra_func
37                    .implicit_params
38                    .iter()
39                    .map(|it| it.to_string())
40                    .collect::<Vec<_>>()
41                    .join(", "),
42                extra_func
43                    .explicit_params
44                    .iter()
45                    .map(|it| it.to_string())
46                    .collect::<Vec<_>>()
47                    .join(", "),
48                extra_func
49            )?;
50        }
51        Ok(())
52    }
53}
54
55/// Debug display for the program state.
56impl Display for Function {
57    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
58        let global_nums = self
59            .analysis_cache
60            .try_get::<GlobalValues>()
61            .unwrap_or_default();
62        let liveness = self
63            .analysis_cache
64            .try_get::<Liveness>()
65            .unwrap_or_else(|| Rc::new(Liveness::empty(self)));
66        let shared_liveness = self
67            .analysis_cache
68            .try_get::<SharedLiveness>()
69            .unwrap_or_else(|| Rc::new(SharedLiveness::empty(self)));
70        let uniformity = self
71            .analysis_cache
72            .try_get::<Uniformity>()
73            .unwrap_or_default();
74
75        if DEBUG_GVN {
76            writeln!(f, "# Value Table:")?;
77            writeln!(f, "{}", global_nums.borrow().values)?;
78        }
79
80        let smems = shared_liveness
81            .allocations
82            .values()
83            .map(|it| format!("    {it}"));
84        let smems = smems.collect::<Vec<_>>().join(",\n");
85        writeln!(f, "Shared memories: [\n{smems}\n]\n")?;
86
87        for node in self.node_indices() {
88            let id = node.index();
89            let bb = &self[node];
90            let uniform = match uniformity.is_block_uniform(node) {
91                true => "uniform ",
92                false => "",
93            };
94            writeln!(f, "{uniform}bb{id} {{")?;
95            if DEBUG_GVN {
96                let block_sets = &global_nums
97                    .borrow()
98                    .block_sets
99                    .get(&node)
100                    .cloned()
101                    .unwrap_or_default();
102                writeln!(f, "{block_sets}")?;
103            }
104
105            if !bb.block_use.is_empty() {
106                writeln!(f, "    Uses: {:?}", bb.block_use)?;
107            }
108            let live_vars = liveness.at_block(node).iter();
109            let live_vars = live_vars.map(|it| format!("%{it}"));
110            let live_vars = live_vars.collect::<Vec<_>>();
111            writeln!(f, "    Live variables: [{}]\n", live_vars.join(", "))?;
112            let live_shared = shared_liveness.at_block(node).iter();
113            let live_shared = live_shared.map(|it| format!("shared({it})"));
114            let live_shared = live_shared.collect::<Vec<_>>();
115            writeln!(
116                f,
117                "    Live shared memories: [{}]\n",
118                live_shared.join(", ")
119            )?;
120
121            for phi in bb.phi_nodes.borrow().iter() {
122                write!(f, "    {} = phi ", phi.out)?;
123                for entry in &phi.entries {
124                    write!(f, "[bb{}: ", entry.block.index())?;
125                    write!(f, "{}]", entry.value)?;
126                }
127                let is_uniform = match uniformity.is_val_uniform(phi.out) {
128                    true => " @ uniform",
129                    false => "",
130                };
131                writeln!(f, ";{is_uniform}")?;
132            }
133            if !bb.phi_nodes.borrow().is_empty() {
134                writeln!(f)?;
135            }
136
137            for op in bb.ops.borrow_mut().values_mut() {
138                let op_fmt = op.to_string();
139                if op_fmt.is_empty() {
140                    continue;
141                }
142
143                let is_uniform = match op.out.is_some_and(|out| uniformity.is_val_uniform(out)) {
144                    true => " @ uniform",
145                    false => "",
146                };
147                writeln!(f, "    {op_fmt};{is_uniform}")?;
148            }
149            match &*bb.control_flow.borrow() {
150                ControlFlow::IfElse {
151                    cond,
152                    then,
153                    or_else,
154                    merge,
155                } => {
156                    writeln!(
157                        f,
158                        "    {cond} ? bb{} : bb{}; merge: {}",
159                        then.index(),
160                        or_else.index(),
161                        merge
162                            .as_ref()
163                            .map(|it| format!("bb{}", it.index()))
164                            .unwrap_or("None".to_string())
165                    )?;
166                }
167                super::ControlFlow::Switch {
168                    value,
169                    default,
170                    branches,
171                    ..
172                } => {
173                    write!(f, "    switch({value}) ")?;
174                    for (val, block) in branches {
175                        write!(f, "[{val}: bb{}] ", block.index())?;
176                    }
177                    writeln!(f, "[default: bb{}];", default.index())?;
178                }
179                super::ControlFlow::Loop {
180                    body,
181                    continue_target,
182                    merge,
183                } => {
184                    writeln!(
185                        f,
186                        "    loop(continue: bb{}, merge: bb{})",
187                        continue_target.index(),
188                        merge.index()
189                    )?;
190                    writeln!(f, "    branch bb{};", body.index())?
191                }
192                super::ControlFlow::LoopBreak {
193                    break_cond,
194                    body,
195                    continue_target,
196                    merge,
197                } => {
198                    writeln!(
199                        f,
200                        "    loop(cond: {}, body: bb{} continue: bb{}, break: bb{})",
201                        break_cond,
202                        body.index(),
203                        continue_target.index(),
204                        merge.index()
205                    )?;
206                }
207                super::ControlFlow::Return { value } => writeln!(
208                    f,
209                    "    return{};",
210                    value.map(|it| format!(" {it}")).unwrap_or_default()
211                )?,
212                super::ControlFlow::Unreachable => writeln!(f, "    unreachable;")?,
213                super::ControlFlow::None => {
214                    let edge = self.edges(node).next();
215                    let target = edge.map(|it| it.target().index()).unwrap_or(255);
216                    writeln!(f, "    branch bb{target};")?;
217                }
218            }
219            f.write_str("}\n\n")?;
220        }
221
222        Ok(())
223    }
224}
225
226impl Display for BlockSets {
227    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
228        let mut exp_gen = self.exp_gen.iter().collect::<Vec<_>>();
229        exp_gen.sort_by_key(|it| it.0);
230        let exp_gen = exp_gen
231            .into_iter()
232            .map(|(val, expr)| format!("{val}: {expr}"))
233            .collect::<Vec<_>>();
234        let mut phi_gen = self.phi_gen.iter().collect::<Vec<_>>();
235        phi_gen.sort_by_key(|it| it.0);
236        let phi_gen = phi_gen
237            .into_iter()
238            .map(|(val, expr)| format!("{val}: {expr}"))
239            .collect::<Vec<_>>();
240        let tmp_gen = self
241            .tmp_gen
242            .iter()
243            .map(|it| format!("{it}"))
244            .collect::<Vec<_>>();
245        let mut leaders = self.leaders.iter().collect::<Vec<_>>();
246        leaders.sort_by_key(|it| it.0);
247        let leaders = leaders
248            .into_iter()
249            .map(|(val, expr)| format!("{val}: {expr}"))
250            .collect::<Vec<_>>();
251        let mut antic_out = self.antic_out.iter().collect::<Vec<_>>();
252        antic_out.sort_by_key(|it| it.0);
253        let antic_out = antic_out
254            .into_iter()
255            .map(|(val, expr)| format!("{val}: {expr}"))
256            .collect::<Vec<_>>();
257        let mut antic_in = self.antic_in.iter().collect::<Vec<_>>();
258        antic_in.sort_by_key(|it| it.0);
259        let antic_in = antic_in
260            .into_iter()
261            .map(|(val, expr)| format!("{val}: {expr}"))
262            .collect::<Vec<_>>();
263
264        writeln!(f, "    exp_gen: [{}]", exp_gen.join(", "))?;
265        writeln!(f, "    phi_gen: [{}]", phi_gen.join(", "))?;
266        writeln!(f, "    tmp_gen: [{}]", tmp_gen.join(", "))?;
267        writeln!(f, "    leaders: [{}]", leaders.join(", "))?;
268        writeln!(f, "    antic_in: [{}]", antic_in.join(", "))?;
269        writeln!(f, "    antic_out: [{}]", antic_out.join(", "))
270    }
271}
272
273impl Display for ValueTable {
274    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
275        let mut values = self.value_numbers.iter().collect::<Vec<_>>();
276        values.sort_by_key(|it| it.1);
277        writeln!(f, "values: [")?;
278        for (val, num) in values {
279            writeln!(f, "    {num}: {val},")?;
280        }
281        writeln!(f, "]")?;
282        writeln!(f, "expressions: [")?;
283        let mut expressions = self.expression_numbers.iter().collect::<Vec<_>>();
284        expressions.sort_by_key(|it| it.1);
285        for (expr, val) in expressions {
286            writeln!(f, "    {val}: {expr},")?;
287        }
288        writeln!(f, "]")
289    }
290}
291
292impl Display for Expression {
293    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
294        match self {
295            Expression::Instruction(instruction) => write!(f, "{instruction}"),
296            Expression::Copy(val, _) => write!(f, "copy({val})"),
297            Expression::Value(value) => write!(f, "{value}"),
298            Expression::Volatile(value) => write!(f, "volatile({value})"),
299            Expression::Phi(entries) => write!(
300                f,
301                "phi({})",
302                entries
303                    .iter()
304                    .map(|(val, b)| format!("{val}: bb{}", b.index()))
305                    .collect::<Vec<_>>()
306                    .join(", ")
307            ),
308            Expression::Builtin(builtin, _) => write!(f, "builtin({builtin:?})"),
309        }
310    }
311}
312
313impl Display for Instruction {
314    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
315        write!(f, "{:?}: [{:?}]", self.op, self.args)
316    }
317}
318
319impl Display for BasicBlock {
320    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
321        for phi in self.phi_nodes.borrow().iter() {
322            write!(f, "    {} = phi ", phi.out)?;
323            for entry in &phi.entries {
324                write!(f, "[bb{}: ", entry.block.index())?;
325                write!(f, "{}]", entry.value)?;
326            }
327            writeln!(f, ";\n")?;
328        }
329        if !self.phi_nodes.borrow().is_empty() {
330            writeln!(f)?;
331        }
332
333        for op in self.ops.borrow_mut().values_mut() {
334            let op_fmt = op.to_string();
335            if op_fmt.is_empty() {
336                continue;
337            }
338
339            writeln!(f, "    {op_fmt};")?;
340        }
341        match &*self.control_flow.borrow() {
342            ControlFlow::IfElse {
343                cond,
344                then,
345                or_else,
346                merge,
347            } => {
348                writeln!(
349                    f,
350                    "    {cond} ? bb{} : bb{}; merge: {}",
351                    then.index(),
352                    or_else.index(),
353                    merge
354                        .as_ref()
355                        .map(|it| format!("bb{}", it.index()))
356                        .unwrap_or("None".to_string())
357                )?;
358            }
359            super::ControlFlow::Switch {
360                value,
361                default,
362                branches,
363                ..
364            } => {
365                write!(f, "    switch({value}) ")?;
366                for (val, block) in branches {
367                    write!(f, "[{val}: bb{}] ", block.index())?;
368                }
369                writeln!(f, "[default: bb{}];", default.index())?;
370            }
371            super::ControlFlow::Loop {
372                body,
373                continue_target,
374                merge,
375            } => {
376                writeln!(
377                    f,
378                    "    loop(continue: bb{}, merge: bb{})",
379                    continue_target.index(),
380                    merge.index()
381                )?;
382                writeln!(f, "    branch bb{};", body.index())?
383            }
384            super::ControlFlow::LoopBreak {
385                break_cond,
386                body,
387                continue_target,
388                merge,
389            } => {
390                writeln!(
391                    f,
392                    "    loop(cond: {}, body: bb{} continue: bb{}, break: bb{})",
393                    break_cond,
394                    body.index(),
395                    continue_target.index(),
396                    merge.index()
397                )?;
398            }
399            super::ControlFlow::Return { value } => writeln!(
400                f,
401                "    return{};",
402                value.map(|it| format!(" {it}")).unwrap_or_default()
403            )?,
404            super::ControlFlow::Unreachable => writeln!(f, "    unreachable;")?,
405            super::ControlFlow::None => {
406                writeln!(f, "    branch;")?;
407            }
408        }
409        Ok(())
410    }
411}
412
413impl Display for SmemAllocation {
414    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
415        let crate::MemoryBlock {
416            value_ty,
417            alignment,
418            ..
419        } = self.smem;
420        write!(
421            f,
422            "shared(id: {}, offset: {}, align: {alignment}, ty: {value_ty})",
423            self.id, self.offset,
424        )
425    }
426}
427
428impl Function {
429    pub fn dot_viz(&self) -> String {
430        let uniformity = self.analysis_cache.try_get::<Uniformity>();
431
432        let get_node_attributes = |_, (index, bb)| {
433            let uniform = uniformity
434                .as_ref()
435                .map(|uniformity| uniformity.is_block_uniform(index))
436                .unwrap_or(false);
437            let title = match uniform {
438                true => format!("uniform bb{}", index.index()),
439                false => format!("bb{}", index.index()),
440            };
441            let bb = format!("{bb}");
442            let lines = bb
443                    .lines()
444                    .map(|it| it.trim())
445                    .map(escape_html)
446                    .filter(|it| !it.is_empty())
447                    .enumerate()
448                    .map(|(i, it)| {
449                        format!(r#"<TR><TD ALIGN="LEFT"><FONT COLOR="dimgray">{i} </FONT></TD><TD ALIGN="LEFT">{it}</TD></TR>"#)
450                    })
451                    .collect::<Vec<_>>();
452            format!(
453                r#"label = <
454<TABLE ALIGN="LEFT" BORDER="0" CELLSPACING="0" CELLBORDER="1" CELLPADDING="0">
455    <TR><TD BGCOLOR="lightgray" ALIGN="LEFT" CELLPADDING="3"><B>{title}</B></TD></TR>
456    <TR><TD CELLPADDING="4">
457    <TABLE ALIGN="LEFT" BORDER="0" CELLSPACING="0" CELLPADDING="0">
458        {}
459    </TABLE>
460    </TD></TR>
461</TABLE>>"#,
462                lines.join("")
463            )
464        };
465
466        //Dot::with_config(&self.program, &[Config::EdgeNoLabel])
467        let content: Dot<'_, &StableDiGraph<BasicBlock, u32>> = Dot::with_attr_getters(
468            &self.graph,
469            &[
470                Config::EdgeNoLabel,
471                Config::NodeNoLabel,
472                Config::GraphContentOnly,
473            ],
474            &|_, _| String::new(),
475            &get_node_attributes,
476        );
477        format!(
478            r#"
479digraph {{
480    node [ shape = box, fontname = "Consolas, 'Courier New', monospace", fontsize = "12", margin = 0 ]
481{content}
482}}
483"#
484        )
485    }
486}
487
488fn escape_html(s: &str) -> String {
489    s.replace("&", "&amp;")
490        .replace("<", "&lt;")
491        .replace(">", "&gt;")
492}
cubecl_opt/debug.rs

cubecl_opt/
debug.rs
