From 48601e2229d27512251e89a7686beb15d8005ca1 Mon Sep 17 00:00:00 2001
From: Russel Arbore <russel.jma@gmail.com>
Date: Wed, 29 Jan 2025 22:15:48 -0600
Subject: [PATCH] Visualize types in dot

---
 hercules_ir/src/dot.rs   | 24 ++++++++++++++++++------
 juno_scheduler/src/pm.rs |  3 +++
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/hercules_ir/src/dot.rs b/hercules_ir/src/dot.rs
index 8efabd7a..5ccda9dc 100644
--- a/hercules_ir/src/dot.rs
+++ b/hercules_ir/src/dot.rs
@@ -16,6 +16,7 @@ use crate::*;
 pub fn xdot_module(
     module: &ir::Module,
     reverse_postorders: &Vec<Vec<NodeID>>,
+    typing: Option<&ModuleTyping>,
     doms: Option<&Vec<DomTree>>,
     fork_join_maps: Option<&Vec<HashMap<NodeID, NodeID>>>,
     devices: Option<&Vec<Device>>,
@@ -30,6 +31,7 @@ pub fn xdot_module(
     write_dot(
         &module,
         &reverse_postorders,
+        typing,
         doms,
         fork_join_maps,
         devices,
@@ -53,6 +55,7 @@ pub fn xdot_module(
 pub fn write_dot<W: Write>(
     module: &ir::Module,
     reverse_postorders: &Vec<Vec<NodeID>>,
+    typing: Option<&ModuleTyping>,
     doms: Option<&Vec<DomTree>>,
     fork_join_maps: Option<&Vec<HashMap<NodeID, NodeID>>>,
     devices: Option<&Vec<Device>>,
@@ -89,6 +92,7 @@ pub fn write_dot<W: Write>(
                 function_id,
                 color,
                 module,
+                typing,
                 &function.schedules[node_id.idx()],
                 w,
             )?;
@@ -249,6 +253,7 @@ fn write_node<W: Write>(
     function_id: FunctionID,
     color: &str,
     module: &Module,
+    typing: Option<&ModuleTyping>,
     schedules: &Vec<Schedule>,
     w: &mut W,
 ) -> std::fmt::Result {
@@ -331,30 +336,37 @@ fn write_node<W: Write>(
     } else {
         format!("{} ({})", node.upper_case_name(), suffix)
     };
+    let xlabel = format!("{}", node_id.idx());
+    let mut tylabel = String::new();
+    if let Some(ty) = typing.map(|typing| typing[function_id.idx()][node_id.idx()]) {
+        module.write_type(ty, &mut tylabel)?;
+    }
 
     let mut iter = schedules.into_iter();
     if let Some(first) = iter.next() {
-        let subtitle = iter.fold(format!("{:?}", first), |b, i| format!("{}, {:?}", b, i));
+        let schedules = iter.fold(format!("{:?}", first), |b, i| format!("{}, {:?}", b, i));
         write!(
             w,
-            "{}_{}_{} [xlabel={}, label=<{}<BR /><FONT POINT-SIZE=\"8\">{}</FONT>>, color={}];\n",
+            "{}_{}_{} [xlabel={}, label=<{}<BR /><FONT POINT-SIZE=\"8\">{}</FONT><BR /><FONT POINT-SIZE=\"8\">{}</FONT>>, color={}];\n",
             node.lower_case_name(),
             function_id.idx(),
             node_id.idx(),
-            node_id.idx(),
+            xlabel,
             label,
-            subtitle,
+            tylabel,
+            schedules,
             color
         )?;
     } else {
         write!(
             w,
-            "{}_{}_{} [xlabel={}, label=\"{}\", color={}];\n",
+            "{}_{}_{} [xlabel={}, label=<{}<BR /><FONT POINT-SIZE=\"8\">{}</FONT>>, color={}];\n",
             node.lower_case_name(),
             function_id.idx(),
             node_id.idx(),
-            node_id.idx(),
+            xlabel,
             label,
+            tylabel,
             color
         )?;
     }
diff --git a/juno_scheduler/src/pm.rs b/juno_scheduler/src/pm.rs
index e62bc78d..c388833c 100644
--- a/juno_scheduler/src/pm.rs
+++ b/juno_scheduler/src/pm.rs
@@ -1636,12 +1636,14 @@ fn run_pass(
 
             pm.make_reverse_postorders();
             if force_analyses {
+                pm.make_typing();
                 pm.make_doms();
                 pm.make_fork_join_maps();
                 pm.make_devices();
             }
 
             let reverse_postorders = pm.reverse_postorders.take().unwrap();
+            let typing = pm.typing.take();
             let doms = pm.doms.take();
             let fork_join_maps = pm.fork_join_maps.take();
             let devices = pm.devices.take();
@@ -1650,6 +1652,7 @@ fn run_pass(
                 xdot_module(
                     module,
                     &reverse_postorders,
+                    typing.as_ref(),
                     doms.as_ref(),
                     fork_join_maps.as_ref(),
                     devices.as_ref(),
-- 
GitLab