graphs of llops, sorry about the line count
MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -8,23 +8,60 @@ ProcessingOps = Enum("ProcessingOps", ["CONV", "CONVT", "CONVDW"])
import os
DEBUG = int(os.getenv("PRINT_LLOPS", "0"))
+GRAPH = int(os.getenv("GRAPH_LLOPS", "0"))
+if GRAPH:
+ import atexit
+ import networkx as nx
+ G = nx.DiGraph()
+ def save_graph_exit():
+ print("saving", G)
+ nx.drawing.nx_pydot.write_dot(G, '/tmp/net.dot')
+ atexit.register(save_graph_exit)
+
+global_num_max = 0
+def log_op(op, ret, inp):
+ if DEBUG: print(f"{op} : {', '.join([str(x.shape) for x in inp])} -> {ret.shape}")
+ if GRAPH:
+ def nm(x):
+ global global_num_max
+ if getattr(x, 'global_num', None) is None:
+ setattr(x, 'global_num', global_num_max)
+ global_num_max += 1
+ return f"<<< {x.global_num} >>>"
+
+ top,sop = str(op).split(".")
+ top_colors = {
+ "UnaryOps": "#c0c0c0",
+ "ReduceOps": "#8080ff",
+ "BinaryOps": "#c0c0c0",
+ "MovementOps": "#80ff80",
+ "ProcessingOps": "#ff8080"
+ }
+
+ for x in inp:
+ G.add_edge(nm(x), nm(ret), label=sop)
+ G.nodes[nm(x)]['label'] = str(x.shape)
+ G.nodes[nm(ret)]['label'] = str(ret.shape)
+ G.nodes[nm(ret)]['fillcolor'] = top_colors[top]
+ G.nodes[nm(ret)]['style'] = 'filled'
+
class Ops:
- def unary_op(ctx, op, x, ret):
- if DEBUG: print(f"{op} : {x.shape} -> {ret.shape}")
+ def unary_op(ctx, op:UnaryOps, x, ret):
+ log_op(op, ret, [x])
return ctx.op.unary_op(op, x, ret)
- def reduce_op(ctx, op, inp, ret):
- if DEBUG: print(f"{op} : {inp.shape} -> {ret.shape}")
- return ctx.op.reduce_op(op, inp, ret)
+ def reduce_op(ctx, op:BinaryOps, x, ret):
+ log_op(op, ret, [x])
+ return ctx.op.reduce_op(op, x, ret)
- def binary_op(ctx, op, x, y, ret):
- if DEBUG: print(f"{op} : {x.shape} + {y.shape} -> {ret.shape}")
+ def binary_op(ctx, op:ReduceOps, x, y, ret):
+ log_op(op, ret, [x, y])
return ctx.op.binary_op(op, x, y, ret)
- def movement_op(ctx, op, inp, ret, arg=None):
- if DEBUG: print(f"{op} : {inp.shape} -> {ret.shape}")
- return ctx.op.movement_op(op, inp, ret, arg)
+ def movement_op(ctx, op:MovementOps, x, ret, arg=None):
+ log_op(op, ret, [x])
+ return ctx.op.movement_op(op, x, ret, arg)
- def processing_op(ctx, op, x, y, ret, stride, groups):
- if DEBUG: print(f"{op} : {x.shape} + {y.shape} -> {ret.shape}")
+ def processing_op(ctx, op:ProcessingOps, x, y, ret, stride, groups):
+ log_op(op, ret, [x, y])
return ctx.op.processing_op(op, x, y, ret, stride, groups)
\ No newline at end of file

graphs of llops, sorry about the line count