diff --git a/plugins/mindstudio-insight-plugins/ModelVis/rust/parser/src/dot/mod.rs b/plugins/mindstudio-insight-plugins/ModelVis/rust/parser/src/dot/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..7be2148ecacea3ecae506b66e68290ee6ec51f2f --- /dev/null +++ b/plugins/mindstudio-insight-plugins/ModelVis/rust/parser/src/dot/mod.rs @@ -0,0 +1,273 @@ +// Copyright (c) 2025, Huawei Technologies Co., Ltd. +// All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + + +use std::{ + fs::File, + io::{BufReader, Read}, +}; + +use ahash::{HashMap, HashMapExt}; +use anyhow::{Result, anyhow}; +use graphviz_rust::dot_structures::{ + Attribute, Edge as DotEdge, EdgeTy, Graph, GraphAttributes, Id, Node as DotNode, Stmt, + Subgraph as DotSubgraph, Vertex, +}; +use smartstring::alias::String; + +use crate::{AttrValue, Model, Node, StdString, Subgraph}; + +pub fn read_dot(path: &str) -> Result { + let file = File::open(path)?; + let mut reader = BufReader::new(file); + let mut content = StdString::new(); + reader.read_to_string(&mut content)?; + + Ok(content) +} + +struct DotParsingContext { + nodes: HashMap, + edges: Vec<(String, String)>, + graph_attributes: HashMap, + subgraphs: HashMap, + nesting_map: HashMap, + current_subgraph: Option, +} + +impl DotParsingContext { + fn new() -> Self { + Self { + nodes: HashMap::new(), + edges: Vec::new(), + graph_attributes: HashMap::new(), + subgraphs: HashMap::new(), + nesting_map: HashMap::new(), + current_subgraph: None, + } + } + + fn into_model(self, name: String) -> Model { + Model { + name, + nodes: self.nodes, + edges: self.edges, + parameters: self.graph_attributes, + subgraphes: self.subgraphs, + nesting_map: self.nesting_map, + } + } + + fn process_statement(&mut self, stmt: &Stmt) -> Result<()> { + match stmt { + Stmt::Subgraph(sg) => self.process_subgraph(sg)?, + Stmt::Node(n) => self.process_node(n)?, + Stmt::Edge(e) => self.process_edge(e)?, + Stmt::GAttribute(ga) => self.process_graph_attribute(ga)?, + Stmt::Attribute(a) => self.process_attribute(a)?, + } + + Ok(()) + } + + fn process_subgraph(&mut self, sg: &DotSubgraph) -> Result<()> { + let name = plain_id(&sg.id)?; + + self.subgraphs + .entry(name.clone()) + .or_insert_with(|| Subgraph { name: name.clone() }); + + let prev_parent = self.current_subgraph.take(); + self.current_subgraph = Some(name); + + for stmt in &sg.stmts { + self.process_statement(stmt)?; + } + + self.current_subgraph = prev_parent; + + Ok(()) + } + + fn process_node(&mut self, n: &DotNode) -> Result<()> { + let node_id = plain_id(&n.id.0)?; + + let mut op_type = node_id.clone(); + let mut attributes = HashMap::new(); + + for attr in &n.attributes { + let key = plain_id(&attr.0)?; + if &key == "pos" { + return Ok(()); + } + + let value = plain_id(&attr.1)?; + + if key == "label" { + op_type = value.clone(); + } + + attributes.insert(key, AttrValue::StringLike(value)); + } + + let node = Node { + name: node_id.clone(), + opType: op_type, + input: Vec::new(), + output: Vec::new(), + attributes, + tensors: Vec::new(), + dynamic: false, + }; + + self.nodes.insert(node_id.clone(), node); + + if let Some(subgraph) = &self.current_subgraph { + self.nesting_map.insert(node_id, subgraph.clone()); + } + + Ok(()) + } + + fn ensure_node_exists(&mut self, node_id: String) { + if !self.nodes.contains_key(&node_id) { + let node = Node { + name: node_id.clone(), + opType: node_id.clone(), + input: Vec::new(), + output: Vec::new(), + attributes: HashMap::new(), + tensors: Vec::new(), + dynamic: false, + }; + self.nodes.insert(node_id.clone(), node); + + if let Some(subgraph) = &self.current_subgraph { + self.nesting_map.insert(node_id, subgraph.clone()); + } + } + } + + fn process_edge(&mut self, e: &DotEdge) -> Result<()> { + match &e.ty { + EdgeTy::Pair(source, target) => match (source, target) { + (Vertex::N(src), Vertex::N(tgt)) => { + let src_id = plain_id(&src.0)?; + let tgt_id = plain_id(&tgt.0)?; + + self.ensure_node_exists(src_id.clone()); + self.ensure_node_exists(tgt_id.clone()); + + self.edges.push((src_id, tgt_id)); + + Ok(()) + } + _ => Err(anyhow!("Subgraph edges are not supported")), + }, + EdgeTy::Chain(vertices) => { + for i in 0..vertices.len() - 1 { + let src = &vertices[i]; + let tgt = &vertices[i + 1]; + + if let (Vertex::N(src), Vertex::N(tgt)) = (src, tgt) { + let src_id = plain_id(&src.0)?; + let tgt_id = plain_id(&tgt.0)?; + + self.ensure_node_exists(src_id.clone()); + self.ensure_node_exists(tgt_id.clone()); + + self.edges.push((src_id, tgt_id)); + } + } + + Ok(()) + } + } + } + + fn process_attribute(&mut self, attr: &Attribute) -> Result<()> { + if self.current_subgraph.is_none() { + let key = plain_id(&attr.0)?; + if &key == "pos" { + return Ok(()); + } + + let value = plain_id(&attr.1)?; + self.graph_attributes.insert(key, value); + } + Ok(()) + } + + fn process_graph_attribute(&mut self, ga: &GraphAttributes) -> Result<()> { + if self.current_subgraph.is_none() { + let attributes = match ga { + GraphAttributes::Graph(attrs) => attrs, + _ => return Ok(()), + }; + + for attr in attributes { + let key = plain_id(&attr.0)?; + if &key == "pos" || &key == "bb" || &key == "size" { + continue; + } + + let value = plain_id(&attr.1)?; + self.graph_attributes.insert(key, value); + } + } + + Ok(()) + } + + fn populate_node_connections(&mut self) { + for (src, tgt) in &self.edges { + if let Some(src_node) = self.nodes.get_mut(src) { + src_node.output.push(tgt.clone()); + } + if let Some(tgt_node) = self.nodes.get_mut(tgt) { + tgt_node.input.push(src.clone()); + } + } + } +} + +fn plain_id(id: &Id) -> Result { + match id { + Id::Plain(p) => Ok(String::from(p)), + Id::Escaped(e) => Ok(String::new()), + _ => Err(anyhow!(format!("Only Plaint Text id is supported, actual: {id}"))), + } +} + +pub fn parse_dot(raw: &str) -> Result { + let graph = graphviz_rust::parse(raw).map_err(|e| anyhow!(e))?; + + match graph { + Graph::Graph { .. } => Err(anyhow!("Undirected graphs are not supported")), + Graph::DiGraph { id, stmts, .. } => { + let mut ctx = DotParsingContext::new(); + + for stmt in &stmts { + ctx.process_statement(stmt)?; + } + + ctx.populate_node_connections(); + + let name = plain_id(&id)?; + + Ok(ctx.into_model(name)) + } + } +} diff --git a/plugins/mindstudio-insight-plugins/ModelVis/rust/parser/src/lib.rs b/plugins/mindstudio-insight-plugins/ModelVis/rust/parser/src/lib.rs index 9061b18a66d049978b1f67cfb3e4ba2eed3a55b4..aa6da760c41516f064a89fea969f2d60888b821a 100644 --- a/plugins/mindstudio-insight-plugins/ModelVis/rust/parser/src/lib.rs +++ b/plugins/mindstudio-insight-plugins/ModelVis/rust/parser/src/lib.rs @@ -27,6 +27,9 @@ use processors::*; mod pbtxt; use pbtxt::*; +mod dot; +use dot::*; + mod dispatch; use dispatch::{FileType, FileType::*}; @@ -36,13 +39,16 @@ pub use str_ext::StrExt; pub mod string_ext; pub use string_ext::SmartStringExt; -mod reader; - pub type StdString = String; pub fn parse_bin(path: &str) -> Result { if path.ends_with(".pbtxt") { return parse_onnx_pbtxt(path) } + if path.ends_with(".dot") { + let content = read_dot(path)?; + return parse_dot(&content) + } + match FileType::from(path) { ONNX => parse_onnx(path), MindIR => parse_mindir(path),