This commit is contained in:
George Hotz
2025-12-21 21:39:17 -04:00
parent 76f2f14233
commit 43c033e5e5

View File

@@ -3,7 +3,7 @@ from collections import defaultdict
from tinygrad.engine.realize import ExecItem
from tinygrad.device import Device, Buffer
from tinygrad.helpers import NO_MEMORY_PLANNER, dedup, DEBUG, round_up
from tinygrad.uop.ops import Ops
from tinygrad.uop.ops import Ops, UOp
from tinygrad.dtype import dtypes, ImageDType
from tinygrad.runtime.support.memory import TLSFAllocator
@@ -63,8 +63,6 @@ def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ign
return assigned
from tinygrad.uop.ops import UOp
def memory_planner(schedule:list[ExecItem]) -> tuple[list[ExecItem], dict[UOp, Buffer]]:
# Exclude buffers involved in load ops (e.g transfers) to preserve parallelism in graphs.
assigned = _internal_memory_planner([[b for b in si.bufs if b is not None] for si in schedule],