github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 23:48:01 -05:00
Code Issues Packages Projects Releases Wiki Activity

am: allow allocations as small as a page (#9523)

Browse Source 
* am: fix allocs

* bettermsg

* comment

* next time
This commit is contained in:
nimlgen
2025-03-21 15:53:32 +08:00
committed by GitHub
parent 3ad3ac4d1e
commit 8a131ab271
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tinygrad/runtime/support/am/amdev.py
View File

@@ -220,13 +220,15 @@ class AMMemoryManager:
if contigous: paddrs = [(self.palloc(size, zero=True), size)]
else:
# Traverse the PT to find the largest contiguous sizes we need to allocate. Try to allocate the longest segment to reduce TLB pressure.
paddrs = []
ctx = AMPageTableTraverseContext(self.adev, self.root_page_table, va, create_pts=True)
for off, _, _, seg_cnt, seg_size in ctx.next(size):
while seg_cnt > 0:
rem_len = seg_cnt * seg_size
while rem_len > 0:
# Try to allocate as long segment (power of 2) as possible
cont_seg_sz, paddr = 1 << (self._frag_size(ctx.vaddr+off, seg_cnt*seg_size) + 12), None
while cont_seg_sz >= seg_size:
cont_seg_sz, paddr = 1 << (self._frag_size(ctx.vaddr+off, rem_len) + 12), None
while cont_seg_sz >= 0x1000:
try: paddr = self.palloc(cont_seg_sz, zero=True)
except MemoryError: cont_seg_sz //= 2
else: break
@@ -234,8 +236,8 @@ class AMMemoryManager:
if paddr is not None: paddrs += [(paddr, cont_seg_sz)]
else:
for paddr, _ in paddrs: self.pa_allocator.free(paddr)
raise MemoryError(f"Failed to allocate contigous {cont_seg_sz=:#x} bytes (size={size:#x})")
seg_cnt, off = seg_cnt - cont_seg_sz // seg_size, off + cont_seg_sz
raise MemoryError(f"Failed to allocate contigous a page. (allocation size={size:#x})")
rem_len, off = rem_len - cont_seg_sz, off + cont_seg_sz
return self.map_range(va, size, paddrs, uncached=uncached)