new memory scheduler with explicit refcounts (#4198)

* new memory scheduler with explict refcounts

* move central memory planner

* typo + use central memory planner in openpilot

* cleanups

* include lb_refcount in pickle

* replace PlaceHolder with memory planner

* cleaner
This commit is contained in:
George Hotz
2024-04-17 08:46:47 +04:00
committed by GitHub
parent c91b7b1739
commit 8564e28a1b
6 changed files with 67 additions and 46 deletions

View File

@@ -14,7 +14,7 @@ from extra.onnx import get_run_onnx
from tinygrad import Tensor, Device, GlobalCounters, dtypes
from tinygrad.dtype import ImageDType
from tinygrad.helpers import partition, Context, fetch, getenv, DEBUG
from tinygrad.engine.realize import run_schedule
from tinygrad.engine.realize import run_schedule, memory_planner
from tinygrad.engine.schedule import create_schedule
from tinygrad.ops import LoadOps, ScheduleItem
Device.DEFAULT = "GPU"
@@ -107,6 +107,7 @@ if __name__ == "__main__":
run_schedule(schedule_independent)
run_schedule(schedule_input)
schedule = memory_planner(schedule)
with Context(DEBUG=max(DEBUG.value, 2), BEAM=getenv("LATEBEAM")):
image_count = sum(isinstance(out.dtype, ImageDType) for si in schedule for out in si.outputs)
print(f"**** running real kernels {image_count}/{len(schedule)} images ****")