nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

memory: use page shifts in memory manager (#11149)

Browse Source 
* memory: use page shifts in memory manager

* fix
This commit is contained in:
nimlgen
2025-07-09 22:05:00 +03:00
committed by GitHub
parent 5c1d215b41
commit b6981404ed
4 changed files with 13 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/external/external_test_am.py
Vendored
+5 -4
View File
@@ -21,12 +21,13 @@ class FakeAM:
def __init__(self):
self.is_booting, self.smi_dev = True, False
self.pcidev = FakePCIDev()
self.vram_mv = memoryview(bytearray(4 << 30))
self.vram_size = (4 << 30)
self.vram_mv = memoryview(bytearray(self.vram_size))
self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
self.gmc = FakeGMC(self)
self.mm = AMMemoryManager(self, 4 << 30, boot_size=(32 << 20), pt_t=AMPageTableEntry, pte_cnt=[512, 512, 512, 512],
pte_covers=[(1 << ((9 * (3-lv)) + 12)) for lv in range(4)], first_lv=am.AMDGPU_VM_PDB1, va_base=AMMemoryManager.va_allocator.base,
palloc_ranges=[(1 << i, 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), 11, -1)])
self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39], va_bits=48,
first_lv=am.AMDGPU_VM_PDB1, va_base=AMMemoryManager.va_allocator.base,
palloc_ranges=[(1 << (i + 12), 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
self.is_booting = False
self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
tinygrad/runtime/support/am/amdev.py
+2 -2
View File
@@ -163,8 +163,8 @@ class AMDev(PCIDevImplBase):
self.smi_dev = smi_dev # During boot only boot memory can be allocated. This flag is to validate this.
# Memory manager & firmware
self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, pte_cnt=[512, 512, 512, 512],
pte_covers=[(1 << ((9 * (3-lv)) + 12)) for lv in range(4)], first_lv=am.AMDGPU_VM_PDB1, va_base=AMMemoryManager.va_allocator.base,
self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39], va_bits=48,
first_lv=am.AMDGPU_VM_PDB1, va_base=AMMemoryManager.va_allocator.base,
palloc_ranges=[(1 << (i + 12), 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
self.fw = AMFirmware(self)
tinygrad/runtime/support/memory.py
+4 -3
View File
@@ -158,10 +158,11 @@ class PageTableTraverseContext:
class MemoryManager:
va_allocator: ClassVar[TLSFAllocator|None] = None
def __init__(self, dev, vram_size:int, boot_size:int, pt_t, pte_cnt:list[int], pte_covers:list[int], va_base:int,
def __init__(self, dev, vram_size:int, boot_size:int, pt_t, va_bits:int, va_shifts:list[int], va_base:int,
palloc_ranges:list[tuple[int, int]], first_lv:int=0):
self.dev, self.vram_size, self.va_base = dev, vram_size, va_base
self.pt_t, self.pte_cnt, self.pte_covers, self.palloc_ranges = pt_t, pte_cnt, pte_covers, palloc_ranges
self.dev, self.vram_size, self.va_shifts, self.va_base, lvl_msb = dev, vram_size, va_shifts, va_base, va_shifts + [va_bits + 1]
self.pte_covers, self.pte_cnt = [1 << x for x in va_shifts][::-1], [1 << (lvl_msb[i+1] - lvl_msb[i]) for i in range(len(lvl_msb) - 1)][::-1]
self.pt_t, self.palloc_ranges = pt_t, palloc_ranges
self.boot_allocator = TLSFAllocator(boot_size, base=0) # per device
self.pa_allocator = TLSFAllocator(vram_size - (64 << 20), base=self.boot_allocator.size) # per device
tinygrad/runtime/support/nv/nvdev.py
+2 -2
View File
@@ -79,8 +79,8 @@ class NVDev(PCIDevImplBase):
# 2 PDE1 (or 512M PTE) 37:29
# 3 PDE0 (dual 64k/4k PDE, or 2M PTE) 28:21
# 4 PTE_64K / PTE_4K 20:16 / 20:12
self.mm = NVMemoryManager(self, self.vram_size, boot_size=(2 << 20), pt_t=NVPageTableEntry, pte_cnt=[4, 512, 512, 256, 512], va_base=0,
pte_covers=[0x800000000000, 0x4000000000, 0x20000000, 0x200000, 0x1000], palloc_ranges=[(x, x) for x in [0x20000000, 0x200000, 0x1000]])
self.mm = NVMemoryManager(self, self.vram_size, boot_size=(2 << 20), pt_t=NVPageTableEntry, va_bits=48, va_shifts=[12, 21, 29, 38, 47], va_base=0,
palloc_ranges=[(x, x) for x in [0x20000000, 0x200000, 0x1000]])
self.flcn:NV_FLCN = NV_FLCN(self)
self.gsp:NV_GSP = NV_GSP(self)