From b6981404ed2bafe9246c4b3da2f455ef9b6ea3c7 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Wed, 9 Jul 2025 22:05:00 +0300 Subject: [PATCH] memory: use page shifts in memory manager (#11149) * memory: use page shifts in memory manager * fix --- test/external/external_test_am.py | 9 +++++---- tinygrad/runtime/support/am/amdev.py | 4 ++-- tinygrad/runtime/support/memory.py | 7 ++++--- tinygrad/runtime/support/nv/nvdev.py | 4 ++-- 4 files changed, 13 insertions(+), 11 deletions(-) diff --git a/test/external/external_test_am.py b/test/external/external_test_am.py index 180c15193f..ca4da691c5 100644 --- a/test/external/external_test_am.py +++ b/test/external/external_test_am.py @@ -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) diff --git a/tinygrad/runtime/support/am/amdev.py b/tinygrad/runtime/support/am/amdev.py index 8a7433388a..7c7b615997 100644 --- a/tinygrad/runtime/support/am/amdev.py +++ b/tinygrad/runtime/support/am/amdev.py @@ -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) diff --git a/tinygrad/runtime/support/memory.py b/tinygrad/runtime/support/memory.py index d24a61b238..8d88edcb21 100644 --- a/tinygrad/runtime/support/memory.py +++ b/tinygrad/runtime/support/memory.py @@ -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 diff --git a/tinygrad/runtime/support/nv/nvdev.py b/tinygrad/runtime/support/nv/nvdev.py index 468255665e..3172f1ba3f 100644 --- a/tinygrad/runtime/support/nv/nvdev.py +++ b/tinygrad/runtime/support/nv/nvdev.py @@ -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)