github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-07 03:00:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

new utils for scheduler graph rewrite [run_process_replay] (#6485)

Browse Source
This commit is contained in:
qazal
2024-09-12 10:01:24 +08:00
committed by GitHub
parent a17ea53340
commit 00d4bf16d8
2 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tinygrad/engine/graph.py
View File

@@ -73,7 +73,8 @@ def log_lazybuffer(lb:'LazyBuffer', scheduled=False):
uops_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.REDUCE: "#C4A484",
UOps.RANGE: "#c8a0e0", UOps.ASSIGN: "#e0ffc0", UOps.BARRIER: "#ff8080", UOps.IF: "#c8b0c0"}
UOps.RANGE: "#c8a0e0", UOps.ASSIGN: "#e0ffc0", UOps.BARRIER: "#ff8080", UOps.IF: "#c8b0c0",
UOps.SWIZZLE: "#7ACD93", UOps.SHAPETRACKER: "#C8F9D4", UOps.REDUCE_AXIS: "#f58488"}
graph_uops_cnt = 0
def word_wrap(x, wrap=80): return x if len(x) <= wrap else (x[0:wrap] + "\n" + word_wrap(x[wrap:], wrap))
def graph_uops(uops:List[UOp]):

7
tinygrad/engine/schedule.py
View File

@@ -179,7 +179,12 @@ def _lower_lazybuffer(outs:List[LazyBuffer], realizes:Dict[LazyBuffer, None]) ->
ubuf = UOp(UOps.DEFINE_GLOBAL, out.dtype if isinstance(out.dtype, ImageDType) else PtrDType(out.dtype), (), i)
ast.append(UOp(UOps.STORE, dtypes.void, (ubuf, output_st.to_uop(), src)))
sink = UOp(UOps.SINK, dtypes.void, tuple(ast))
if AST_REWRITE: sink = graph_rewrite(sink, reduceop_fusor)
if AST_REWRITE:
try: sink = graph_rewrite(sink, reduceop_fusor)
except Exception as e:
from tinygrad.engine.graph import graph_uop
graph_uop(sink)
raise e
return LBScheduleItem(sink, outs, list(inputs), var_vals, dedup([x[0].metadata for x in cache if x[0].metadata and x[0] not in inputs]))
# *** DAG creation: decide which LazyBuffers should realize ***