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

make schedule explicit about kernels [pr] (#10462)

Browse Source
This commit is contained in:
George Hotz
2025-05-22 09:32:16 -07:00
committed by GitHub
parent c46edbf262
commit ab591fa4dd
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tinygrad/engine/schedule.py
View File

@@ -40,13 +40,17 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
children: defaultdict[UOp, list[UOp]] = defaultdict(list)
in_degree: dict[UOp, int] = {}
for u in (toposort:=sched_sink.toposort()):
if u.op is not Ops.ASSIGN: continue
if u.op is not Ops.ASSIGN: continue # anything that's not an ASSIGN doesn't write a kernel, so we can skip
k = u.src[1]
in_degree.setdefault(k, 0)
for s in k.src:
if s.op is not Ops.ASSIGN: continue
children[s.src[1]].append(k)
in_degree[k] += 1
if s.op is Ops.ASSIGN:
children[s.src[1]].append(k)
in_degree[k] += 1
elif s.op is Ops.BUFFER:
pass # a BUFFER is already realized, nothing to do here
else:
raise RuntimeError(f"input to kernel must be ASSIGN or BUFFER, not {s.op}")
# linearize KERNEL UOps into ScheduleItems in BFS order
queue = deque(k for k,v in in_degree.items() if v == 0)