github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

memplan: fix large buffers planning (#11878)

Browse Source 
* memplan: fix large buffers planning

* fix

* fix dsp
This commit is contained in:
nimlgen
2025-08-27 23:54:27 +03:00
committed by GitHub
parent 4006366752
commit 44816218b5
2 changed files with 16 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
test/test_memory_planner.py
View File

@@ -120,5 +120,19 @@ 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()

3
tinygrad/engine/memory.py
View File

@@ -23,12 +23,13 @@ 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()) * 2 # *2 for fragmentation (which is about 15%)
# 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(1 << 44, block_size=0x1000, lv2_cnt=32)))
global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=min_block_size, 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):