[Visualizer] update logic (#121)

* gitignore

* init

* [Draft]Complete basic structure, still met with type conversion problem

* Update core.pyx

* Visualizer python interface

* support newly added optypes

* Restore cutlass submodule to commit cc3c29a

* Repair minor issues including replace magic type numbers, useless var in func and remove block_dim

* fit in kn output node

---------

Co-authored-by: Mengdi Wu <[email protected]>
Co-authored-by: Jianan Ji <[email protected]>

python/mirage/visualizer.py

@@ -21,7 +21,7 @@
         "bg": "#fdf2d0",
         "edge": "#b89230",
         "edge_label": "black",
-        "io": "#4273b1",  # 输入输出节点和边的颜色
+        "io": "#4273b1",
     }
 }
 
@@ -131,8 +131,9 @@ def is_kernel_output_node(self):
         return not self.output_tensors
 
     def draw(self, G):
-        G.node(self.name, label=self.label, color=self.color, style="rounded,filled", shape="box",
-               penwidth="0", fontsize=node_font_size, fontcolor="white", fontname="sans-serif", margin="0.4,0.4")
+        if not self.is_output_node():
+            G.node(self.name, label=self.label, color=self.color, style="rounded,filled", shape="box",
+                penwidth="0", fontsize=node_font_size, fontcolor="white", fontname="sans-serif", margin="0.4,0.4")
 
 class block_node(node):
     def __init__(self, name, op_type, id, label, iomap=None, forloop_dim=None, forloop_range=None):
@@ -215,35 +216,23 @@ def draw(self, G):
                         penwidth="0", fontsize=tensor_node_font_size, fontcolor="white", fontname="sans-serif")
                     if not (self.last_node.is_input_node()):
                         draw_edge(G, self.last_node.name, self.name, "kernel")
-                    if next_node:
+                    if not (next_node.is_output_node()):
                         draw_edge(G, self.name, next_node.name, "kernel")
                 elif self.last_node.is_input_node():
                     if not self.name:
                         self.name = str(self.guid)
                     G.node(self.name, label=str(self.shape), color=self.color, style="filled", shape="box",
                         penwidth="0", fontsize=tensor_node_font_size, fontcolor="white", fontname="sans-serif")
-                    draw_edge(G, self.name, next_node.name, "kernel")
-                elif not next_node:
+                    if not (next_node.is_output_node()):
+                        draw_edge(G, self.name, next_node.name, "kernel")
+                elif next_node.is_output_node():
                     if not self.name:
                         self.name = str(self.guid)
                     G.node(self.name, label=str(self.shape), color=self.color, style="filled", shape="box",
                         penwidth="0", fontsize=tensor_node_font_size, fontcolor="white", fontname="sans-serif")
                     draw_edge(G, self.last_node.name, self.name, "kernel")
                 else:
                     draw_edge(G, self.last_node.name, next_node.name, "kernel", label=str(self.shape))
-        elif self.last_node.is_customized_node():
-                if not self.name:
-                    G.node(self.name, label=str(self.shape), color=self.color, style="filled", shape="box",
-                    penwidth="0", fontsize=tensor_node_font_size, fontcolor="white", fontname="sans-serif")
-                else:
-                    G.node(self.name, label=self.name + "\n" + str(self.shape), color=self.color, style="filled", shape="box",
-                        penwidth="0", fontsize=tensor_node_font_size, fontcolor="white", fontname="sans-serif")
-                draw_edge(G, self.last_node.name, self.name, "kernel")
-        else:
-            G.node(self.last_node.name+"_output_tensor", label=str(self.shape), color=self.color, style="filled", shape="box",
-                penwidth="0", fontsize=tensor_node_font_size, fontcolor="white", fontname="sans-serif")
-            draw_edge(G, self.last_node.name, self.last_node.name+"_output_tensor", "kernel")
-
 class block_tensor(tensor):
     def __init__(self, guid, shape):
         super()._init_(guid, colors_map["block"]["edge"], shape)
@@ -449,7 +438,6 @@ def handle_graph_data(graph_data, graph_title, output_filename, dot=True, png=Tr
 
 
 if __name__ == "__main__":
-    # file_name = "multi_graph.json" 
     file_name = "mirage_search_checkpoint.json"
     with open(file_name) as f:
         data = json.load(f)