tract_libcli/
export.rs

1use crate::annotations::Annotations;
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
29#[derive(Clone, Debug, Serialize)]
30pub struct ProfilingInfo {
31    iterations: usize,
32    secs_per_iter: f64,
33}
34
35impl GraphPerfInfo {
36    pub fn from(model: &dyn Model, annotations: &Annotations) -> GraphPerfInfo {
37        let nodes = annotations
38            .tags
39            .iter()
40            .map(|(id, node)| Node {
41                qualified_id: NodeQIdSer(id.0.iter().cloned().collect(), id.1),
42                cost: node.cost.iter().map(|(k, v)| (format!("{k:?}"), format!("{v}"))).collect(),
43                node_name: id.model(model).unwrap().node_name(id.1).to_string(),
44                op_name: id.model(model).unwrap().node_op_name(id.1).to_string(),
45                secs_per_iter: node.profile.map(|s| s.as_secs_f64()),
46            })
47            .collect();
48        let profiling_info = annotations.profile_summary.as_ref().map(|summary| ProfilingInfo {
49            secs_per_iter: summary.entire.as_secs_f64(),
50            iterations: summary.iters,
51        });
52        GraphPerfInfo { nodes, profiling_info }
53    }
54}