early src delete (#3996)

* early src delete

* fix bad test

* fix test_linearizer
This commit is contained in:
George Hotz
2024-03-29 19:46:07 -07:00
committed by GitHub
parent 31c8ba8b84
commit 8f1e34a2a0
3 changed files with 10 additions and 5 deletions

View File

@@ -44,7 +44,6 @@ def run_schedule(schedule:List[ScheduleItem]):
for out in si.outputs:
# we don't have an output buffer, we have to create it, and create to max size if it has symbolic shape
if out.size > 0 and not dont_allocate and out.op is not LoadOps.ASSIGN: out.buffer.allocate()
del out.srcs
# run the function (put it in JIT)
real_buffers = [x.buffer for x in si.outputs+si.inputs if x.size != 0]

View File

@@ -197,6 +197,7 @@ def create_schedule(outs:List[LazyBuffer], seen:Optional[Set[LazyBuffer]]=None)
graph[out].append(assign_targets[x])
in_degree[assign_targets[x]] += 1
if x in prescheduled: in_degree[out] += 1
del out.srcs # can only schedule once
queue = deque(out for out in prescheduled if in_degree[out] == 0)
schedule: List[ScheduleItem] = []