update llm.c with UOp ast [run_process_replay] (#6296)

examples/llm.c/export.py

@@ -7,7 +7,7 @@ from train_gpt2 import GPT, GPTConfig
 from tinygrad.helpers import dedup, to_function_name, flatten, getenv, GRAPH, GlobalCounters, ansilen, to_function_name
 from tinygrad.engine.schedule import create_schedule
 from tinygrad.engine.realize import get_kernel, memory_planner, run_schedule
-from tinygrad.ops import BufferOps, MetaOps
+from tinygrad.ops import MetaOps, UOps
 
 TIMING = getenv("TIMING")
 
@@ -43,7 +43,7 @@ if __name__ == "__main__":
     #run_schedule(sched[:])
   del seen  # free the LazyBuffers
   sched = memory_planner(sched)
-  ast_dedup = dedup([si.ast for si in sched if si.ast[0].op is BufferOps.STORE])
+  ast_dedup = dedup([si.ast for si in sched if si.ast.op is UOps.SINK])
   srcs = {}
   for ast in ast_dedup:
     k = get_kernel(Device["CLANG"].renderer, ast)
@@ -84,8 +84,8 @@ if __name__ == "__main__":
   for i,si in enumerate(sched):
     bufs = [(named_buffers.get(b, f"b{numbered_bufs[b]}"), b) for b in si.bufs]
     all_bufs += bufs
-    if si.ast[0].op is not BufferOps.STORE:
-      print(f"// {si.ast[0].op}", bufs)
+    if si.ast.op is not UOps.SINK:
+      print(f"// {si.ast.op}", bufs)
     else:
       print(f"{srcs[si.ast][0]}({', '.join([x[0] for x in bufs])})")
       main.append(f"  {to_function_name(srcs[si.ast][0])}({', '.join([x[0] for x in bufs])});")