From 2e6162b2817d5d6a024e30b5e15c8f0dc04dee12 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Tue, 16 Jan 2024 20:57:31 -0800
Subject: [PATCH] graph cleanup (#3155)

* simpler graph

* unused functions
---
 tinygrad/graph.py | 36 +++++++++++++-----------------------
 1 file changed, 13 insertions(+), 23 deletions(-)

diff --git a/tinygrad/graph.py b/tinygrad/graph.py
index 1c45cb0b74..dcf28d92ef 100644
--- a/tinygrad/graph.py
+++ b/tinygrad/graph.py
@@ -1,12 +1,15 @@
-import os, atexit
+import os, atexit, functools
 from collections import defaultdict
 from typing import List, Any, DefaultDict
-from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, BufferOps, TernaryOps, Op, LazyOp, GlobalCounters
+from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, BufferOps, TernaryOps, LazyOp, GlobalCounters
 from tinygrad.device import Device
 from tinygrad.helpers import GRAPHPATH, DEBUG, getenv
 from tinygrad.codegen.linearizer import UOps, UOp
 from tinygrad.shape.symbolic import NumNode
 
+try: import networkx as nx
+except ImportError: pass
+
 # **** debugging and graphing ****
 
 if DEBUG >= 2:
@@ -16,18 +19,17 @@ if DEBUG >= 2:
           f"{' '*10}total: {GlobalCounters.kernel_count:5d} kernels {GlobalCounters.global_ops*1e-9:8.2f} GOPS {GlobalCounters.global_mem*1e-9:8.2f} GB {GlobalCounters.time_sum_s*1e3:8.2f} ms")  # noqa: E501
   atexit.register(print_globalcounters)
 
+def save_graph(G, fn, opt=""):
+  print("saving", G, f"to {fn}.svg")
+  nx.drawing.nx_pydot.write_dot(G, f'{fn}.dot')
+  os.system(f'dot {opt} -Tsvg {fn}.dot -o {fn}.svg')
+
 G:Any = None
 def init_graph():
   global G
   if G is not None: return
-  import networkx as nx
   G = nx.DiGraph()
-  def save_graph_exit():
-    print("saving", G, f"to {GRAPHPATH}.svg")
-    nx.drawing.nx_pydot.write_dot(G, f'{GRAPHPATH}.dot')
-    # -Gnslimit=100 can make it finish, but you won't like results
-    os.system(f'dot -Tsvg {GRAPHPATH}.dot -o {GRAPHPATH}.svg')
-  atexit.register(save_graph_exit)
+  atexit.register(functools.partial(save_graph, G, GRAPHPATH)) # -Gnslimit=100 can make it finish, but you won't like results
 
 counts: DefaultDict[type, int] = defaultdict(int)
 def nm(x):
@@ -36,16 +38,6 @@ def nm(x):
     counts[type(x)] += 1
   return x.node_id
 
-def get_sop(op: List[Op]):
-  op = [x for x in op if x not in BufferOps]
-  if len(op) <= 2: return '.'.join([str(y).split(".")[1] for y in op][::-1])
-  if len(op) <= 6: return '.'.join([str(y).split(".")[1][0:3] for y in op][::-1])
-  return str(len(op))
-
-def str_dtype(dtyp):
-  ret = str(dtyp)[7:]
-  return "" if ret == 'float' else f"\n{ret}"
-
 def realized_lazybuffer(lb, num):
   init_graph()
   G.nodes[nm(lb)]['style'] = '"filled,bold"'
@@ -68,7 +60,7 @@ def log_lazybuffer(lb, scheduled=False):
       G.add_edge(nm(x), nm(lb), color='#a0a0a0')
     label = '"' + \
       (str(set(x.shape for x in lb.srcs))+"\n"+str(lb.shape) if lb.op in ReduceOps else str(lb.shape)) + \
-      str_dtype(lb.dtype)+f"\n{lb.op}"+(f"\n{lb.arg}" if lb.op in {LoadOps.CONST, UnaryOps.CAST} else "") + \
+      (f"\n{lb.dtype.name}" if lb.dtype.name != "float" else "")+f"\n{lb.op}"+(f"\n{lb.arg}" if lb.op in {LoadOps.CONST, UnaryOps.CAST} else "") + \
       (f"\n{lb.device}" if lb.device != Device.DEFAULT else "") + '"'
     G.add_node(nm(lb), style='"filled,dashed"', fillcolor=[v for k,v in top_colors.items() if lb.op in k][0] + "80", color="black", label=label)
     if scheduled: G.nodes[nm(lb)]['shape'] = 'box'
@@ -91,7 +83,6 @@ def _tree(lazyop:LazyOp, cycles, cnt, prefix=""):
 def print_tree(lazyop:LazyOp): print("\n".join([f"{str(i).rjust(3)} {s}" for i,s in enumerate(_tree(lazyop, {}, [-1]))]))
 
 def graph_uops(uops:List[UOp]):
-  import networkx as nx
   colors = {UOps.ALU: "#ffffc0", UOps.LOAD: "#ffc0c0", UOps.STORE: "#c0ffc0", UOps.SPECIAL: "#c0c0ff", UOps.CONST: "#e0e0e0",
             UOps.DEFINE_GLOBAL: "#ffe0b0", UOps.DEFINE_LOCAL: "#ffe0d0", UOps.DEFINE_ACC: "#f0ffe0",
             UOps.LOOP: "#c8a0e0", UOps.PHI: "#e0ffc0", UOps.BARRIER: "#ff8080", UOps.IF: "#c8b0c0"}
@@ -100,5 +91,4 @@ def graph_uops(uops:List[UOp]):
     if u.uop == UOps.END: continue
     G.add_node(uops.index(u), label=f"{str(u.uop)[5:]}{(' '+str(u.arg)) if u.arg is not None else ''}\n{str(u.dtype)}", style="filled", fillcolor=colors.get(u.uop, "#ffffff"))  # noqa: E501
     for v in u.vin: G.add_edge(uops.index(v), uops.index(u))
-  nx.drawing.nx_pydot.write_dot(G, f'{GRAPHPATH}.uops.dot')
-  os.system(f'dot -Grankdir=LR -Tsvg {GRAPHPATH}.uops.dot -o {GRAPHPATH}.uops.svg')
+  save_graph(G, f'{GRAPHPATH}.uops', '-Grankdir=LR')