From fafbf3daea2213723235a9fe48e4524d5f28dfd3 Mon Sep 17 00:00:00 2001
From: nimlgen <138685161+nimlgen@users.noreply.github.com>
Date: Wed, 15 Oct 2025 22:47:50 +0800
Subject: [PATCH] memory: reserve ptable (#12702)

---
 tinygrad/runtime/support/memory.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tinygrad/runtime/support/memory.py b/tinygrad/runtime/support/memory.py
index 653a397980..c6718fdba3 100644
--- a/tinygrad/runtime/support/memory.py
+++ b/tinygrad/runtime/support/memory.py
@@ -135,7 +135,7 @@ class PageTableTraverseContext:
   def _try_free_pt(self) -> bool:
     pt, _, _ = self.pt_stack[-1]
     if self.free_pts and pt != self.dev.mm.root_page_table and all(not pt.valid(i) for i in range(self._pt_pte_cnt(self.pt_stack[-1][0].lv))):
-      self.dev.mm.pfree(pt.paddr)
+      self.dev.mm.pfree(pt.paddr, ptable=True)
       parent_pt, parent_pte_idx, _ = self.pt_stack[-2]
       parent_pt.set_entry(parent_pte_idx, 0x0, valid=False)
       return True
@@ -258,4 +258,4 @@ class MemoryManager:
     if zero: self.dev.vram[paddr:paddr+size] = bytes(size)
     return paddr
 
-  def pfree(self, paddr:int): self.pa_allocator.free(paddr)
+  def pfree(self, paddr:int, ptable=False): (self.ptable_allocator if self.reserve_ptable and ptable else self.pa_allocator).free(paddr)