From 400636675279e218269da693fe509d626409787d Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Wed, 27 Aug 2025 22:36:14 +0300 Subject: [PATCH] Revert "memplan: fix large buffers planning (#11876)" (#11877) This reverts commit 7f90497efcb5f40f13c88b0bd6c61975945ee137. --- test/test_memory_planner.py | 14 -------------- tinygrad/engine/memory.py | 3 +-- 2 files changed, 1 insertion(+), 16 deletions(-) diff --git a/test/test_memory_planner.py b/test/test_memory_planner.py index 6aa03c6be3..0ba3f90404 100644 --- a/test/test_memory_planner.py +++ b/test/test_memory_planner.py @@ -120,19 +120,5 @@ class TestMemoryPlanner(unittest.TestCase): ] check_assign(bs) - def test_very_small_buffers(self): - bs = [ - [b(0, pin=True), b(1, size=32)], - [b(3, size=4), b(4, size=6)], - ] - check_assign(bs) - - def test_very_big_buffers(self): - bs = [ - [b(0, pin=True), b(1, size=34359738368000)], - [b(3, size=1 << 128), b(4, size=1 << 64)], - ] - check_assign(bs) - if __name__ == "__main__": unittest.main() diff --git a/tinygrad/engine/memory.py b/tinygrad/engine/memory.py index b7b28fc11b..2c62b3f6a1 100644 --- a/tinygrad/engine/memory.py +++ b/tinygrad/engine/memory.py @@ -23,13 +23,12 @@ def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ign # Sort buffer operations in timeline order. Two events: buffer is allocated or buffer is freed. buffer_requests = sorted([((first_appearance[buf], True), buf) for buf in first_appearance.keys()] + \ [((last_appearance[buf] + 1, False), buf) for buf in first_appearance.keys()], key=lambda x: x[0]) - total_memory = sum(round_up(buf.nbytes, min_block_size:=0x1000) for buf in first_appearance.keys()) # Try to suballocate from a shared buffer managed by global_planner using TLSFAllocator. # Also track buffer replacements for buffers that do not support suballocation. buffer_replace:dict[Buffer, tuple[Buffer|None, int|None]] = {} reuse_buffers:dict[tuple, list[Buffer]] = defaultdict(list) - global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=min_block_size, lv2_cnt=32))) + global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(1 << 44, block_size=0x1000, lv2_cnt=32))) for (_, is_open_ev), buf in buffer_requests: # Check if suballocation is possible for the given buffer and device. if hasattr(Device[buf.device].allocator, "_offset") and not isinstance(buf.dtype, ImageDType):