tract_libcli/
export.rs

1use crate::annotations::{Annotations, NodeQId};
2use crate::model::Model;
3use serde::Serialize;
4use std::collections::HashMap;
5use tract_core::internal::*;
6
7#[derive(Clone, Debug, Default, Serialize)]
8pub struct GraphPerfInfo {
9    nodes: Vec<Node>,
10    profiling_info: Option<ProfilingInfo>,
11}
12
13#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize)]
14pub struct NodeQIdSer(pub Vec<(usize, String)>, pub usize);
15
16#[derive(Clone, Debug, Serialize)]
17pub struct Node {
18    qualified_id: NodeQIdSer,
19    op_name: String,
20    node_name: String,
21
22    #[serde(skip_serializing_if = "HashMap::is_empty")]
23    cost: HashMap<String, String>,
24
25    #[serde(skip_serializing_if = "Option::is_none")]
26    secs_per_iter: Option<f64>,
27
28    /// Accelerator time attributed to this node, when the runtime reports it.
29    #[serde(skip_serializing_if = "Option::is_none")]
30    accel_secs_per_iter: Option<f64>,
31}
32
33#[derive(Clone, Debug, Serialize)]
34pub struct ProfilingInfo {
35    iterations: usize,
36    secs_per_iter: f64,
37}
38
39impl GraphPerfInfo {
40    pub fn from(model: &dyn Model, annotations: &Annotations) -> GraphPerfInfo {
41        let nodes = annotations
42            .tags
43            .iter()
44            .map(|(id, node)| Node {
45                qualified_id: NodeQIdSer(id.0.iter().cloned().collect(), id.1),
46                cost: node.cost.iter().map(|(k, v)| (format!("{k:?}"), format!("{v}"))).collect(),
47                node_name: id.model(model).unwrap().node_name(id.1).to_string(),
48                op_name: id.model(model).unwrap().node_op_name(id.1).to_string(),
49                secs_per_iter: node.profile.map(|s| s.as_secs_f64()),
50                accel_secs_per_iter: node.accelerator_profile.map(|s| s.as_secs_f64()),
51            })
52            .collect();
53        let profiling_info = annotations.profile_summary.as_ref().map(|summary| ProfilingInfo {
54            secs_per_iter: summary.entire.as_secs_f64(),
55            iterations: summary.iters,
56        });
57        GraphPerfInfo { nodes, profiling_info }
58    }
59}
60
61// -- audit-json --
62
63#[derive(Serialize)]
64pub struct AuditModel {
65    properties: HashMap<String, String>,
66    assertions: Vec<String>,
67    inputs: Vec<AuditModelIo>,
68    outputs: Vec<AuditModelIo>,
69    nodes: Vec<AuditNode>,
70}
71
72#[derive(Serialize)]
73struct AuditModelIo {
74    name: String,
75    node: usize,
76    slot: usize,
77    fact: String,
78}
79
80#[derive(Serialize)]
81struct AuditNode {
82    id: usize,
83    name: String,
84    op: String,
85    #[serde(skip_serializing_if = "Vec::is_empty")]
86    info: Vec<String>,
87    inputs: Vec<AuditOutletRef>,
88    outputs: Vec<AuditNodeOutput>,
89    #[serde(skip_serializing_if = "HashMap::is_empty")]
90    cost: HashMap<String, String>,
91}
92
93#[derive(Serialize)]
94struct AuditOutletRef {
95    node: usize,
96    slot: usize,
97}
98
99#[derive(Serialize)]
100struct AuditNodeOutput {
101    fact: String,
102    successors: Vec<AuditInletRef>,
103}
104
105#[derive(Serialize)]
106struct AuditInletRef {
107    node: usize,
108    slot: usize,
109}
110
111pub fn audit_json(
112    model: &dyn Model,
113    annotations: &Annotations,
114    writer: impl std::io::Write,
115) -> TractResult<()> {
116    let properties: HashMap<String, String> =
117        model.properties().iter().map(|(k, v)| (k.clone(), format!("{v:?}"))).collect();
118
119    let scope = model.symbols();
120    let assertions: Vec<String> = scope.all_assertions().iter().map(|a| format!("{a}")).collect();
121
122    let inputs: Vec<AuditModelIo> = model
123        .input_outlets()
124        .iter()
125        .map(|o| {
126            Ok(AuditModelIo {
127                name: model.node_name(o.node).to_string(),
128                node: o.node,
129                slot: o.slot,
130                fact: model.outlet_fact_format(*o),
131            })
132        })
133        .collect::<TractResult<_>>()?;
134
135    let outputs: Vec<AuditModelIo> = model
136        .output_outlets()
137        .iter()
138        .map(|o| {
139            Ok(AuditModelIo {
140                name: model.node_name(o.node).to_string(),
141                node: o.node,
142                slot: o.slot,
143                fact: model.outlet_fact_format(*o),
144            })
145        })
146        .collect::<TractResult<_>>()?;
147
148    let nodes: Vec<AuditNode> = (0..model.nodes_len())
149        .map(|id| {
150            let op = model.node_op(id);
151            let info = op.info().unwrap_or_default();
152            let node_inputs: Vec<AuditOutletRef> = model
153                .node_inputs(id)
154                .iter()
155                .map(|o| AuditOutletRef { node: o.node, slot: o.slot })
156                .collect();
157            let node_outputs: Vec<AuditNodeOutput> = (0..model.node_output_count(id))
158                .map(|slot| {
159                    let outlet = OutletId::new(id, slot);
160                    let fact = model.outlet_fact_format(outlet);
161                    let successors: Vec<AuditInletRef> = model
162                        .outlet_successors(outlet)
163                        .iter()
164                        .map(|inlet| AuditInletRef { node: inlet.node, slot: inlet.slot })
165                        .collect();
166                    AuditNodeOutput { fact, successors }
167                })
168                .collect();
169            let cost: HashMap<String, String> = annotations
170                .tags
171                .get(&NodeQId(tvec!(), id))
172                .map(|tags| {
173                    tags.cost.iter().map(|(k, v)| (format!("{k:?}"), format!("{v}"))).collect()
174                })
175                .unwrap_or_default();
176            AuditNode {
177                id,
178                name: model.node_name(id).to_string(),
179                op: model.node_op_name(id).to_string(),
180                info,
181                inputs: node_inputs,
182                outputs: node_outputs,
183                cost,
184            }
185        })
186        .collect();
187
188    let audit = AuditModel { properties, assertions, inputs, outputs, nodes };
189    serde_json::to_writer_pretty(writer, &audit)?;
190    Ok(())
191}
tract_libcli/export.rs

tract_libcli/
export.rs
