github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

pre_schedule itself is complete

Browse Source
This commit is contained in:
George Hotz
2025-12-19 12:54:01 -04:00
parent 3a51104a73
commit 9af2409fc3
2 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tinygrad/engine/jit2.py
View File

@@ -30,6 +30,7 @@ class TinyJit(Generic[ReturnType]):
# process schedule caches to be symbolic
for k,v in input_state_dict.items(): print(k, v.uop.base)
for input_buffers, sched_cache_key in self.schedule_caches:
pre_schedule, combined_sink = schedule_cache[sched_cache_key]
for k,v in input_buffers.items():
print(k.pyrender())
print(v.pyrender())

2
tinygrad/engine/schedule.py
View File

@@ -89,7 +89,7 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ScheduleItem], UOp]:
else:
ast, buf_uops, metadata, fixedvars, bound_ranges = si
fixedvars = fixedvars | {s.src[0].arg[0]:in_ranges[s.src[1]] for s in bound_ranges}
pre_schedule.append(ScheduleItem(ast, (), metadata, fixedvars))
pre_schedule.append(ScheduleItem(ast, buf_uops, metadata, fixedvars))
buf_uops_list.append(UOp.sink(*buf_uops))
sched_ptr += 1
return pre_schedule, UOp.sink(*buf_uops_list)