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):