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

* add a graph_rewrite pass for creating asts [pr]

* disk

* benchmark

Author: Qazalin

tinygrad/engine/schedule.py

@@ -369,8 +369,8 @@ def check_load_st(glbl:UOp, view:UOp):
   (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