From d43485ae9e7cd29439b1ccac2f5bd7e14e72dd3c Mon Sep 17 00:00:00 2001 From: Yingbo Ma Date: Mon, 27 Nov 2023 21:42:48 -0500 Subject: [PATCH] Fix `graph_uops` (#2457) * Load networkx when we need to graph uops * Document GRAPHUOPS * import nx in `graph_uops` --- docs/env_vars.md | 1 + tinygrad/graph.py | 1 + 2 files changed, 2 insertions(+) diff --git a/docs/env_vars.md b/docs/env_vars.md index badb99cf37..1a1039b6ba 100644 --- a/docs/env_vars.md +++ b/docs/env_vars.md @@ -44,6 +44,7 @@ LLVMOPT | [1] | enable slightly more expensive LLVM optimizat OPT | [1-3] | optimization level BEAM | [#] | number of beams in kernel beam search GRAPH | [1] | create a graph of all operations (requires graphviz) +GRAPHUOPS | [1] | create a graph of uops (requires graphviz and saves at /tmp/uops.{svg,dot}) GRAPHPATH | [/path/to] | where to put the generated graph IMAGE | [1] | enable 2d specific optimizations FLOAT16 | [1] | use float16 for images instead of float32 diff --git a/tinygrad/graph.py b/tinygrad/graph.py index 59895a2e0e..98ae8aa227 100644 --- a/tinygrad/graph.py +++ b/tinygrad/graph.py @@ -104,6 +104,7 @@ def _tree(lazydata, prefix=""): def print_tree(lazydata:LazyOp): print("\n".join([f"{str(i).rjust(3)} {s}" for i,s in enumerate(_tree(lazydata))])) 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"}