use std::collections::HashMap;
use super::core::ExecutionGraph;
use crate::brick::exec_graph::node::{EdgeType, ExecutionNode, TransferDirection};
fn build_children_map(
edges: &[crate::brick::exec_graph::node::ExecutionEdge],
node_count: usize,
) -> (HashMap<u32, Vec<u32>>, Vec<u32>) {
let mut children_map: HashMap<u32, Vec<u32>> = HashMap::new();
let mut has_parent: std::collections::HashSet<u32> = std::collections::HashSet::new();
for edge in edges {
if edge.edge_type == EdgeType::Contains || edge.edge_type == EdgeType::Launches {
children_map.entry(edge.src.0).or_default().push(edge.dst.0);
has_parent.insert(edge.dst.0);
}
}
let root_ids: Vec<u32> = (0..node_count as u32).filter(|id| !has_parent.contains(id)).collect();
(children_map, root_ids)
}
fn format_ascii_node(node: &ExecutionNode) -> (String, String) {
match node {
ExecutionNode::Layer { index } => (format!("Layer {}", index), String::new()),
ExecutionNode::Brick { id: brick_id, timing_ns, elements } => (
brick_id.name().to_string(),
format!(" {:.1}µs ({} elem)", *timing_ns as f64 / 1000.0, elements),
),
ExecutionNode::Kernel { name, grid, block, shared_mem, .. } => (
name.clone(),
format!(" <<<{},{},{}>>> smem={}B", grid.0, block.0, block.1, shared_mem),
),
ExecutionNode::Function { name, file, line } => {
let loc = match (file, line) {
(Some(f), Some(l)) => format!(" ({}:{})", f, l),
(None, _) | (_, None) => String::new(),
};
(format!("{}{}", name, loc), String::new())
}
ExecutionNode::Transfer { src, dst, bytes, direction, timing_ns } => {
let timing_str =
timing_ns.map(|ns| format!(" {:.1}µs", ns as f64 / 1000.0)).unwrap_or_default();
(format!("{:?}: {} → {}", direction, src, dst), format!(" {}B{}", bytes, timing_str))
}
ExecutionNode::AsyncTask { name, poll_count, yield_count, total_poll_ns } => {
let efficiency = if *poll_count > 0 { 100.0 / *poll_count as f64 } else { 0.0 };
(
name.clone(),
format!(
" polls:{} yields:{} {:.1}µs ({:.0}% eff)",
poll_count,
yield_count,
*total_poll_ns as f64 / 1000.0,
efficiency
),
)
}
}
}
fn build_ascii_tree(
graph: &ExecutionGraph,
id: u32,
children_map: &HashMap<u32, Vec<u32>>,
prefix: &str,
connector: &str,
output: &mut String,
) {
let (label, info) = format_ascii_node(&graph.nodes[id as usize]);
output.push_str(&format!("{}{}{}{}\n", prefix, connector, label, info));
if let Some(child_ids) = children_map.get(&id) {
let child_count = child_ids.len();
for (i, &child_id) in child_ids.iter().enumerate() {
let is_last = i == child_count - 1;
let new_connector = if is_last { "└── " } else { "├── " };
let new_prefix = if connector.is_empty() {
prefix.to_string()
} else if connector == "└── " {
format!("{} ", prefix)
} else {
format!("{}│ ", prefix)
};
build_ascii_tree(graph, child_id, children_map, &new_prefix, new_connector, output);
}
}
}
fn node_to_dot_label(node: &ExecutionNode) -> (String, &'static str) {
match node {
ExecutionNode::Layer { index } => {
(format!("Layer {}", index), "style=filled,fillcolor=lightblue")
}
ExecutionNode::Brick { id, timing_ns, .. } => (
format!("{}\\n{:.1}µs", id.name(), *timing_ns as f64 / 1000.0),
"style=filled,fillcolor=lightgreen",
),
ExecutionNode::Kernel { name, grid, block, .. } => (
format!("{}\\n<<<{},{},{}>>>", name, grid.0, block.0, block.1),
"style=filled,fillcolor=lightyellow",
),
ExecutionNode::Function { name, file, line } => {
let loc = match (file, line) {
(Some(f), Some(l)) => format!("\\n{}:{}", f, l),
(None, _) | (_, None) => String::new(),
};
(format!("{}{}", name, loc), "style=filled,fillcolor=lightgray")
}
ExecutionNode::Transfer { src, dst, bytes, direction, .. } => {
let dir = match direction {
TransferDirection::H2D => "H2D",
TransferDirection::D2H => "D2H",
TransferDirection::D2D => "D2D",
};
(
format!("{}\\n{}->{}\\n{:.1}MB", dir, src, dst, *bytes as f64 / 1e6),
"style=filled,fillcolor=lightsalmon",
)
}
ExecutionNode::AsyncTask { name, poll_count, yield_count, total_poll_ns } => {
let efficiency = if *poll_count > 0 { 100.0 / *poll_count as f64 } else { 0.0 };
(
format!(
"{}\\npolls:{} yields:{}\\n{:.1}µs ({:.0}%)",
name,
poll_count,
yield_count,
*total_poll_ns as f64 / 1000.0,
efficiency
),
"style=filled,fillcolor=lightcyan",
)
}
}
}
fn edge_to_dot_style(edge_type: &EdgeType) -> &'static str {
match edge_type {
EdgeType::Calls => "style=solid",
EdgeType::Contains => "style=dashed",
EdgeType::Launches => "style=bold,color=red",
EdgeType::Sequence => "style=dotted",
EdgeType::DependsOn => "style=solid,color=blue",
EdgeType::Transfer { .. } => "style=bold,color=orange",
}
}
impl ExecutionGraph {
pub fn to_dot(&self) -> String {
debug_assert!(
self.nodes.len() <= u32::MAX as usize,
"CB-BUDGET: graph node count {} exceeds u32 range",
self.nodes.len()
);
let mut dot = String::from("digraph ExecutionGraph {\n");
dot.push_str(" rankdir=TB;\n");
dot.push_str(" node [shape=box];\n\n");
for (i, node) in self.nodes.iter().enumerate() {
let (label, style) = node_to_dot_label(node);
dot.push_str(&format!(" n{} [label=\"{}\",{}];\n", i, label, style));
}
dot.push('\n');
for edge in &self.edges {
let style = edge_to_dot_style(&edge.edge_type);
dot.push_str(&format!(" n{} -> n{} [{}];\n", edge.src.0, edge.dst.0, style));
}
dot.push_str("}\n");
dot
}
#[cfg(feature = "execution-graph")]
pub fn to_csr(&self) -> trueno_graph::CsrGraph {
use trueno_graph::{CsrGraph, NodeId};
let edges: Vec<(NodeId, NodeId, f32)> =
self.edges.iter().map(|e| (NodeId(e.src.0), NodeId(e.dst.0), e.weight)).collect();
let mut graph = CsrGraph::from_edge_list(&edges).unwrap_or_default();
for (i, node) in self.nodes.iter().enumerate() {
graph.set_node_name(NodeId(i as u32), node.name());
}
graph
}
#[cfg(feature = "presentar-tui")]
pub fn to_tree_node(&self) -> presentar_terminal::TreeNode {
use presentar_terminal::{Color, TreeNode};
let layer_color = Color::new(0.4, 0.6, 1.0, 1.0);
let brick_color = Color::new(0.4, 0.8, 0.4, 1.0);
let kernel_color = Color::new(1.0, 0.8, 0.3, 1.0);
let func_color = Color::new(0.7, 0.7, 0.7, 1.0);
let (children_map, root_ids) = build_children_map(&self.edges, self.nodes.len());
fn build_node(
graph: &ExecutionGraph,
id: u32,
children_map: &HashMap<u32, Vec<u32>>,
layer_color: Color,
brick_color: Color,
kernel_color: Color,
func_color: Color,
) -> TreeNode {
let node = &graph.nodes[id as usize];
let (label, info, color) = match node {
ExecutionNode::Layer { index } => (format!("Layer {}", index), None, layer_color),
ExecutionNode::Brick { id: brick_id, timing_ns, elements } => (
brick_id.name().to_string(),
Some(format!("{:.1}µs ({} elem)", *timing_ns as f64 / 1000.0, elements)),
brick_color,
),
ExecutionNode::Kernel { name, grid, block, shared_mem, .. } => (
name.clone(),
Some(format!("<<<{},{},{}>>> smem={}B", grid.0, block.0, block.1, shared_mem)),
kernel_color,
),
ExecutionNode::Function { name, file, line } => {
let loc = match (file, line) {
(Some(f), Some(l)) => format!(" ({}:{})", f, l),
(None, _) | (_, None) => String::new(),
};
(format!("{}{}", name, loc), None, func_color)
}
ExecutionNode::Transfer { src, dst, bytes, direction, timing_ns } => {
let timing_str = timing_ns
.map(|ns| format!(" {:.1}µs", ns as f64 / 1000.0))
.unwrap_or_default();
(
format!("{:?}: {} → {}", direction, src, dst),
Some(format!("{}B{}", bytes, timing_str)),
Color::new(0.8, 0.4, 0.8, 1.0), )
}
ExecutionNode::AsyncTask { name, poll_count, yield_count, total_poll_ns } => {
let efficiency = if *poll_count > 0 { 100.0 / *poll_count as f64 } else { 0.0 };
(
name.clone(),
Some(format!(
"polls:{} yields:{} {:.1}µs ({:.0}% eff)",
poll_count,
yield_count,
*total_poll_ns as f64 / 1000.0,
efficiency
)),
Color::new(0.4, 0.8, 0.8, 1.0), )
}
};
let mut tree_node = TreeNode::new(id as u64, label).with_color(color);
if let Some(info_str) = info {
tree_node = tree_node.with_info(info_str);
}
if let Some(child_ids) = children_map.get(&id) {
for &child_id in child_ids {
let child = build_node(
graph,
child_id,
children_map,
layer_color,
brick_color,
kernel_color,
func_color,
);
tree_node = tree_node.with_child(child);
}
}
tree_node
}
if root_ids.is_empty() {
TreeNode::new(0, "Empty Graph")
} else if root_ids.len() == 1 {
build_node(
self,
root_ids[0],
&children_map,
layer_color,
brick_color,
kernel_color,
func_color,
)
} else {
let mut root = TreeNode::new(u64::MAX, "Execution Graph")
.with_color(Color::new(0.9, 0.9, 0.9, 1.0));
for &root_id in &root_ids {
let child = build_node(
self,
root_id,
&children_map,
layer_color,
brick_color,
kernel_color,
func_color,
);
root = root.with_child(child);
}
root
}
}
#[must_use]
pub fn to_ascii_tree(&self) -> String {
let (children_map, root_ids) = build_children_map(&self.edges, self.nodes.len());
let mut output = String::new();
if root_ids.is_empty() {
output.push_str("(empty graph)\n");
} else if root_ids.len() == 1 {
build_ascii_tree(self, root_ids[0], &children_map, "", "", &mut output);
} else {
output.push_str("Execution Graph\n");
let root_count = root_ids.len();
for (i, &root_id) in root_ids.iter().enumerate() {
let is_last = i == root_count - 1;
let connector = if is_last { "└── " } else { "├── " };
build_ascii_tree(self, root_id, &children_map, "", connector, &mut output);
}
}
if output.ends_with('\n') {
output.pop();
}
output
}
}