tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 15:36:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

add a graph_rewrite pass for creating asts [pr] (#9765)

Browse Source 
* add a graph_rewrite pass for creating asts [pr]

* disk

* benchmark
This commit is contained in:
qazal
2025-04-07 16:32:11 +08:00
committed by GitHub
parent 07eea567d4
commit 6306dea6e2
1 changed files with 11 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/engine/schedule.py
+11 -8
View File
@@ -369,8 +369,8 @@ fix_kernel_ops = PatternMatcher([
(UPat(Ops.LOAD, src=(UPat.var("glbl"), UPat.var("view"))), check_load_st),
])
def fix_kernel_ast(k:UOp, var_vals:dict[Variable, int]) -> UOp:
assert k.op is Ops.KERNEL, f"kernel isn't kernel, it's {k}"
def fix_kernel_ast(ctx:dict[Variable, int], k:UOp) -> UOp|None:
if k.arg.ast.op in GroupOp.Meta or all(s.op is Ops.STORE for s in k.arg.ast.src): return None
# substitute kernel sources for the target buffer + apply reshapes
parents_rep: dict[UOp, UOp] = {}
for s in k.src:
@@ -380,11 +380,10 @@ def fix_kernel_ast(k:UOp, var_vals:dict[Variable, int]) -> UOp:
# push views to edges
ast = graph_rewrite(graph_rewrite(ast, view_left), view_right)
# add buffer ops + fix_kernel_ops
ast = graph_rewrite(ast, merge_views+add_buffer_ops+fix_kernel_ops, ctx=(var_vals, bufs:=tuple(s.buf_uop for s in k.src)), bottom_up=True)
ast = graph_rewrite(ast, merge_views+add_buffer_ops+fix_kernel_ops, ctx=(ctx, bufs:=tuple(s.buf_uop for s in k.src)), bottom_up=True)
if ast.op is Ops.SINK and not all_same(dev:=[x.device for x in bufs]): raise RuntimeError(f"all buffers must be on the same device: {dev}")
# create subbuffer (TODO: this does not belong here)
if ast.op is Ops.BUFFER_VIEW: buffers[bufs[0]] = (base:=bufs[1].buffer).view(ast.size, ast.dtype, ast.arg[1]*base.dtype.itemsize)
return k.replace(arg=Kernel(ast, k.arg.metadata))
create_ast = PatternMatcher([(UPat(Ops.KERNEL, name="k"), fix_kernel_ast),])
PROCESS_REPLAY_CAPTURE:dict[str, bytes] = {}
if CAPTURE_PROCESS_REPLAY:
@@ -449,6 +448,10 @@ def create_schedule_with_vars(big_sink:UOp) -> tuple[list[ScheduleItem], dict[Va
if getenv("VIZ"): graph_rewrite(sched_sink, PatternMatcher([]), name="View Kernel Graph")
if getenv("VIZ"): graph_rewrite(sched_sink, PatternMatcher([]), name="View Memory Graph")
# unbind var_vals and fix kernel ast
var_vals: dict[Variable, int] = {}
sched_sink = graph_rewrite(sched_sink, create_ast, ctx=var_vals, bottom_up=True)
# final toposort (bfs)
children: dict[UOp, list[UOp]] = {}
in_degree: dict[UOp, int] = {}
@@ -462,11 +465,11 @@ def create_schedule_with_vars(big_sink:UOp) -> tuple[list[ScheduleItem], dict[Va
queue = deque(k for k,v in in_degree.items() if v == 0)
schedule: list[ScheduleItem] = []
var_vals: dict[Variable, int] = {}
while queue:
u = queue.popleft()
# TODO: move this to create_kernels
k = fix_kernel_ast(u.src[1], var_vals)
# map the BUFFER UOp to a subbuffer if it's a BUFFER_VIEW
if (k:=u.src[1]).arg.ast.op is Ops.BUFFER_VIEW:
buffers[k.src[0]] = (base:=k.src[1].buf_uop.buffer).view(k.size, k.arg.ast.dtype, k.arg.ast.arg[1]*base.dtype.itemsize)
schedule.append(ScheduleItem(k.arg.ast, tuple(s.buf_uop.buffer for s in k.src), k.arg.metadata))
for x in children.get(u, []):
in_degree[x] -= 1