1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use std::fmt::Write as _;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;
use std::process::Command;

use crate::graph::{Graph, Operation};

/// Render the given graph as an svg file.
///
/// This assumes that graphviz is installed and available on the path as `dot`.
pub fn graph_to_svg(path: impl AsRef<Path>, graph: &Graph, hide_const: bool, show_ids: bool) -> std::io::Result<()> {
    let path = path.as_ref();

    let path_gv = path.with_extension("gv");
    let path_svg = path.with_extension("svg");

    let output = BufWriter::new(File::create(&path_gv)?);
    graph_to_dot(output, graph, hide_const, show_ids)?;

    let result = Command::new("dot")
        .arg("-Tsvg")
        .arg(path_gv)
        .arg("-o")
        .arg(path_svg)
        .status()?;
    assert!(result.success(), "Running 'dot' failed with status {:?}", result);

    Ok(())
}

/// Render the given graph as a graphviz string.
///
/// This makes no assumptions about the environment.
pub fn graph_to_dot(mut f: impl Write, graph: &Graph, hide_const: bool, show_ids: bool) -> std::io::Result<()> {
    writeln!(f, "digraph {{")?;
    writeln!(f)?;

    for value in graph.values() {
        if hide_const && graph.is_const(value) {
            continue;
        }

        let info = &graph[value];

        let (color, op, attrs_operation) = match info.operation {
            Operation::Input { index } => ("gray", "Input", vec![("index", format!("{}", index))]),
            Operation::Constant { ref tensor } => {
                let mut attrs = vec![];
                if tensor.len() == 1 {
                    attrs.push(("value", format!("{:?}", tensor)));
                }
                ("gray", "Constant", attrs)
            }
            Operation::View { input: _ } => ("brown", "View", vec![]),
            Operation::Broadcast { input: _ } => ("brown", "Broadcast", vec![]),
            Operation::Permute {
                input: _,
                ref permutation,
            } => {
                let attrs = vec![("Permute", format!("{:?}", permutation))];
                ("brown", "permute", attrs)
            }
            Operation::Slice { input: _, axis, range } => {
                let attrs = vec![("axis", format!("{}", axis)), ("range", format!("{}", range))];
                ("brown", "Slice", attrs)
            }
            Operation::Flip { input: _, axis } => ("brown", "Flip", vec![("axis", format!("{}", axis))]),
            Operation::Gather {
                input: _,
                axis,
                indices: _,
            } => ("yellow", "Gather", vec![("axis", format!("{}", axis))]),
            Operation::Concat { inputs: _, axis } => ("yellow", "Concat", vec![("axis", format!("{}", axis))]),
            Operation::Conv {
                input: _,
                filter: _,
                details,
            } => {
                let mut attrs = vec![("kernel", format!("{}x{}", details.kernel_h, details.kernel_w))];
                if details.has_stride() {
                    attrs.push(("stride", format!("{}x{}", details.stride_y, details.stride_x)));
                }
                if !details.keeps_spatial_shape() {
                    attrs.push(("padding", format!("{}x{}", details.padding_y, details.padding_x)));
                }
                ("blue", "Conv", attrs)
            }
            Operation::MatMul { left: _, right: _ } => ("blue", "MatMul", vec![]),
            Operation::Unary { input: _, op } => ("green", "Unary", vec![("op", format!("{:?}", op))]),
            Operation::Binary { left: _, right: _, op } => ("green", "Binary", vec![("op", format!("{:?}", op))]),
            Operation::Softmax { input: _, axis } => ("purple", "Softmax", vec![("axis", format!("{}", axis))]),
            Operation::Layernorm { input: _, axis, eps: _ } => {
                ("purple", "Layernorm", vec![("axis", format!("{}", axis))])
            }
            Operation::Reduce { input: _, ref axes, op } => (
                "purple",
                "Reduce",
                vec![("op", format!("{:?}", op)), ("axes", format!("{:?}", axes))],
            ),
        };

        let mut attrs_general = vec![];
        attrs_general.push(("shape", format!("{}", info.shape)));
        if let Some(output_index) = graph.outputs().iter().position(|&v| v == value) {
            attrs_general.push(("output", format!("{}", output_index)));
        }

        if show_ids {
            let debug_id = &graph[value].debug_id;
            if !debug_id.is_empty() {
                attrs_general.push(("debug_id", format!("{:?}", debug_id)));
            }
        }

        let mut attrs = attrs_general;
        attrs.extend(attrs_operation.into_iter());

        let mut table = String::new();
        writeln!(&mut table, "<TABLE BORDER=\"0\">").unwrap();
        writeln!(&mut table, "<TR><TD>{:?}</TD><TD><B>{}</B></TD></TR>", value, op).unwrap();
        for (key, value) in attrs {
            writeln!(&mut table, "<TR><TD>{}</TD><TD>{}</TD></TR>", key, value).unwrap();
        }
        writeln!(&mut table, "</TABLE>").unwrap();

        let label = table;
        writeln!(
            f,
            "{} [label=<{}>, color={:?}, shape=box, width=2]",
            value.index(),
            label,
            color,
        )?;
    }

    writeln!(f)?;

    for value in graph.values() {
        for operand in graph[value].operation.inputs() {
            if hide_const && graph.is_const(operand) {
                continue;
            }

            writeln!(f, "{} -> {}", operand.index(), value.index())?;
        }
    }

    writeln!(f)?;
    writeln!(f, "}}")?;
    Ok(())
}