From 552fbd5527825a2036a4e0d60777e48f94331f9b Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Tue, 27 Aug 2024 20:04:54 +0800 Subject: [PATCH] update llm.c with UOp ast [run_process_replay] (#6296) --- examples/llm.c/export.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/llm.c/export.py b/examples/llm.c/export.py index 776a28b8a7..7fa1cc1dcb 100755 --- a/examples/llm.c/export.py +++ b/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])});")