tinygrad 0.10.2__py3-none-any.whl → 0.11.0__py3-none-any.whl

tinygrad/runtime/support/ib.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+import resource, ctypes, weakref, functools, itertools, tinygrad.runtime.autogen.ib as ib
+from typing import Iterator
+from dataclasses import dataclass
+from weakref import WeakKeyDictionary
+from tinygrad.device import Buffer, DMACPURef, DMAFdRef
+from tinygrad.helpers import getenv, round_up, DEBUG
+
+DEFAULT_PORT, DEFAULT_GID = getenv("DEFAULT_PORT", 1), getenv("DEFAULT_GID", 3) # DEFAULT_GID=0 for RXE
+IOVA_ALIGN = resource.getpagesize()
+
+def checkz(x, ret=None):
+  assert x == 0, f'{x} != 0 (errno {ctypes.get_errno()})'
+  return ret
+
+@dataclass(frozen=True)
+class SGE:
+  dst_iova: int
+  dst_key: int
+  src_iova: int
+  src_key: int
+  size: int
+
+class IBCtx:
+  def __init__(self, idx:int):
+    # Open the device (aka Host Channel Adapter in ib-speak)
+    devs = ib.ibv_get_device_list(ctypes.byref(ndevs:=ctypes.c_int32()))
+    if idx >= ndevs.value: raise IndexError(f"{idx} > {ndevs.value}")
+    self.ctx = ib.ibv_open_device(devs[idx])
+    ib.ibv_free_device_list(devs)
+
+    # HACK: remove this (and all usage of `ctx.contents.ops`) when clang2py can deal with `static inline` wrapper-functions
+    self.vctx = ctypes.cast(ctypes.addressof(self.ctx.contents) - ib.struct_verbs_context.context.offset, ctypes.POINTER(ib.struct_verbs_context))
+
+    # Get attributes. Something like port_attr.max_msg_sz sound like it might requre taking the min of host's and remote's attributes if they differ
+    self.device_attr = checkz(ib.ibv_query_device(self.ctx, ctypes.byref(da:=ib.struct_ibv_device_attr())), da)
+    self.port_attr = checkz(self.vctx.contents.query_port(self.ctx, DEFAULT_PORT, ctypes.byref(pa:=ib.struct_ibv_port_attr()), ctypes.sizeof(pa)), pa)
+    self.gid_attr = checkz(ib.ibv_query_gid(self.ctx, DEFAULT_PORT, DEFAULT_GID, ctypes.byref(ga:=ib.union_ibv_gid())), ga)
+
+    # Allocate protection domain
+    self.pd = ib.ibv_alloc_pd(self.ctx)
+    self.next_iova: int = IOVA_ALIGN # don't start at zero (nullptr)
+
+    # weakref(buf) => (iova, mr, mr_dealloc). mr_dealloc is kept here to avoid double freeing mrs that are deallocated in __del__
+    self.mrs: WeakKeyDictionary[Buffer, tuple[int, ctypes._Pointer[ib.struct_ibv_mr], weakref.finalize]] = WeakKeyDictionary()
+
+    # Default soft fd limit is 1024, which is not enough, set soft to hard (maximum allowed by the os)
+    IBCtx.rlimit_fix()
+
+  def __del__(self):
+    # must deallocate all mrs in protection domain before deallocating the protection domain
+    if hasattr(self, "mrs"): [fin() for _,_,fin in self.mrs.values()]
+    if hasattr(self, "pd"): ib.ibv_dealloc_pd(self.pd)
+    if hasattr(self, "ctx"): ib.ibv_close_device(self.ctx)
+
+  @functools.cache # run once
+  @staticmethod
+  def rlimit_fix():
+    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+    resource.setrlimit(resource.RLIMIT_NOFILE, (hard, hard))
+    if DEBUG>=2: print(f"IB: Increased fd limit from {soft} to {hard}")
+
+  def alloc_iova(self, size:int, required_offset:int):
+    iova = round_up(self.next_iova - required_offset, IOVA_ALIGN) + required_offset
+    self.next_iova = iova + size
+    return iova
+
+  def reg(self, buf:Buffer) -> tuple[int, ctypes._Pointer[ib.struct_ibv_mr]]:
+    buf = buf.base
+    if buf not in self.mrs:
+      if buf.nbytes > self.device_attr.max_mr_size: raise RuntimeError(f"Buffer too big: {buf.nbytes:#x} > {self.device_attr.max_mr_size:#x}")
+      if len(self.mrs) >= self.device_attr.max_mr: raise RuntimeError(f"Out of memory region cap: {len(self.mrs)} >= {self.device_attr.max_mr}")
+      # Local read is implied (but still have to create the memory region, except for short sends/writes with IBV_SEND_INLINE that are inlined by cpu)
+      mr_flags = ib.IBV_ACCESS_LOCAL_WRITE | ib.IBV_ACCESS_REMOTE_READ | ib.IBV_ACCESS_REMOTE_WRITE
+      match (dmaref:=buf.as_dmaref()):
+        case DMACPURef():
+          iova = self.alloc_iova(dmaref.size, dmaref.addr % IOVA_ALIGN)
+          mr = ib.ibv_reg_mr_iova2(self.pd, ctypes.c_void_p(dmaref.addr), dmaref.size, iova, mr_flags)
+        case DMAFdRef():
+          iova = self.alloc_iova(dmaref.size, dmaref.offset % IOVA_ALIGN)
+          mr = ib.ibv_reg_dmabuf_mr(self.pd, dmaref.offset, dmaref.size, iova, dmaref.fd, mr_flags)
+        case _: raise RuntimeError(f"Unknown type of dma ref: {dmaref}")
+      if not mr: raise RuntimeError(f"Couldn't register memory region for {buf} {dmaref} (errno={ctypes.get_errno()})")
+      self.mrs[buf] = (iova, mr, weakref.finalize(buf, ib.ibv_dereg_mr, mr))
+    return self.mrs[buf][0:2]
+
+class IBConn:
+  def __init__(self, ctx:IBCtx):
+    self.ctx = ctx
+
+    # Create Completion Channel. It is a file descriptor that kernel sends notifications through, not a thing in infiniband spec, just linux-ism
+    self.comp_channel = ib.ibv_create_comp_channel(self.ctx.ctx)
+    # Create Completion Queue. When a Work Request with signaled flag is completed a Completion Queue Entry is pushed onto this queue
+    self.cq = ib.ibv_create_cq(self.ctx.ctx, _capacity:=256, _cq_context:=None, self.comp_channel, _comp_vector:=0)
+    self.pending_wrids: set[int] = set()
+    self.wrid_num: Iterator[int] = itertools.count(0) # wc_id is uint64, this will never overflow
+
+    # Create Queue Pair. It's the closest thing to a socket in infiniband with QP num being the closest thing to a port, except it's allocated by hca
+    qp_init_attrs_cap = ib.struct_ibv_qp_cap(max_send_wr=1024, max_recv_wr=64, max_send_sge=8, max_recv_sge=8, max_inline_data=64)
+    qp_init_attrs = ib.struct_ibv_qp_init_attr(send_cq=self.cq, recv_cq=self.cq, cap=qp_init_attrs_cap, qp_type=ib.IBV_QPT_RC) # Reliable Connection
+    self.qp = ib.ibv_create_qp(self.ctx.pd, ctypes.byref(qp_init_attrs))
+    self.qp_cap = qp_init_attrs.cap
+
+    # The most important thing about QPs is their state, when a new QP is created it's in the RESET state, before it can be properly used it has to go
+    # through Init, Ready To Receive, Ready To Send. A good docs on QP state machine: https://www.rdmamojo.com/2012/05/05/qp-state-machine/
+
+    # INIT
+    qp_access_flags = ib.IBV_ACCESS_REMOTE_WRITE | ib.IBV_ACCESS_REMOTE_READ
+    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_INIT, port_num=DEFAULT_PORT, qp_access_flags=qp_access_flags)
+    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_PORT | ib.IBV_QP_ACCESS_FLAGS | ib.IBV_QP_PKEY_INDEX))
+
+    self.gid, self.qp_num = bytes(self.ctx.gid_attr.raw), self.qp.contents.qp_num
+
+  # Exchange GID and QP num with remote. At least in RoCEv2 gid can be guessed from remote's ip, QP num can't.
+
+  def connect(self, remote_gid:bytes, remote_qp_num:int):
+    # RTR
+    qp_ah_attr_grh = ib.struct_ibv_global_route(hop_limit=1, dgid=ib.union_ibv_gid(raw=(ctypes.c_ubyte * 16)(*remote_gid)), sgid_index=DEFAULT_GID)
+    qp_ah_attr = ib.struct_ibv_ah_attr(is_global=1, port_num=DEFAULT_PORT, grh=qp_ah_attr_grh)
+    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_RTR, path_mtu=ib.IBV_MTU_4096, dest_qp_num=remote_qp_num, rq_psn=0, max_dest_rd_atomic=1,
+                                min_rnr_timer=12, ah_attr=qp_ah_attr)
+    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_PATH_MTU | ib.IBV_QP_DEST_QPN | ib.IBV_QP_RQ_PSN | \
+                                          ib.IBV_QP_MAX_DEST_RD_ATOMIC | ib.IBV_QP_MIN_RNR_TIMER | ib.IBV_QP_AV))
+
+    # RTS
+    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_RTS, timeout=14, retry_cnt=7, rnr_retry=7, sq_psn=0, max_rd_atomic=1)
+    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_TIMEOUT | ib.IBV_QP_RETRY_CNT | ib.IBV_QP_RNR_RETRY | ib.IBV_QP_SQ_PSN | \
+                                          ib.IBV_QP_MAX_QP_RD_ATOMIC))
+
+  def __del__(self):
+    self.wait_cq() # need to wait for **everything** to complete before it's safe to dealloc queues and stuff
+    ib.ibv_destroy_qp(self.qp)
+    ib.ibv_destroy_cq(self.cq)
+    ib.ibv_destroy_comp_channel(self.comp_channel)
+
+  def next_wrid(self):
+    self.pending_wrids.add(wrid:=next(self.wrid_num))
+    return wrid
+
+  def wait_cq(self, wr_id: int|None=None):
+    while (wr_id in self.pending_wrids) if wr_id is not None else self.pending_wrids:
+      if self.ctx.ctx.contents.ops.poll_cq(self.cq, _num_entries:=1, ctypes.byref(wc:=ib.struct_ibv_wc())):
+        if wc.status != ib.IBV_WC_SUCCESS:
+          raise RuntimeError(f'Work Request completed with error: wr_id={wc.wr_id} status={ib.ibv_wc_status__enumvalues.get(wc.status, wc.status)}')
+        self.pending_wrids.remove(wc.wr_id)
+
+  def rdma_write(self, sgl:list[SGE]):
+    swr: ctypes._Pointer[ib.struct_ibv_send_wr]|None = None
+    swr_cnt, wr_id = 0, self.next_wrid()
+    def _post():
+      nonlocal swr, swr_cnt, wr_id
+      if swr is not None:
+        # The swr can be freed when this returns, the memory that sge points to can be unmapped after work completion is retrieved from cq
+        checkz(self.ctx.ctx.contents.ops.post_send(self.qp, swr, ctypes.byref(_bad_wr:=ctypes.POINTER(ib.struct_ibv_send_wr)())))
+        # TODO: async
+        self.wait_cq(wr_id)
+        swr, swr_cnt, wr_id = None, 0, self.next_wrid()
+    # Everything is in reverse for elegant chaining
+    for sg in reversed(sgl):
+      # Message size limit (max 2GB per ib spec, 1GB on tinybox mellanoxes) applies to both scatter-gather entries and entire wrs
+      for off in reversed(range(0, sg.size, self.ctx.port_attr.max_msg_sz)):
+        # Scatter-Gather Entry for local memory
+        sge = ctypes.pointer(ib.struct_ibv_sge(addr=sg.src_iova+off, length=min(sg.size-off, self.ctx.port_attr.max_msg_sz), lkey=sg.src_key))
+        # RDMA struct for remote memory
+        wr = ib.union_ibv_send_wr_wr(rdma=ib.struct_ibv_send_wr_1_rdma(remote_addr=sg.dst_iova+off, rkey=sg.dst_key))
+        # Signal (with chosen work request id) if it's the last wr (first in the loop since it's reversed)
+        wid, flags = (wr_id, ib.IBV_SEND_SIGNALED) if swr is None else (0, 0)
+        # Create Send Request
+        swr = ctypes.pointer(ib.struct_ibv_send_wr(opcode=ib.IBV_WR_RDMA_WRITE, sg_list=sge, num_sge=1, wr=wr, wr_id=wid, send_flags=flags, next=swr))
+        # Flush if queue is being overrun
+        if (swr_cnt:=swr_cnt + 1) >= self.qp_cap.max_send_wr: _post()
+    _post()

tinygrad/runtime/support/llvm.py

@@ -9,15 +9,14 @@ if sys.platform == 'win32':
     raise FileNotFoundError('LLVM not found, you can install it with `winget install LLVM.LLVM` or point at a custom dll with LLVM_PATH')
 elif OSX:
   # Will raise FileNotFoundError if brew is not installed
-  brew_prefix = subprocess.check_output(['brew', '--prefix', 'llvm']).decode().strip()
   # `brew --prefix` will return even if formula is not installed
-  if not os.path.exists(brew_prefix):
-    raise FileNotFoundError('LLVM not found, you can install it with `brew install llvm`')
+  if not os.path.exists(brew_prefix:=subprocess.check_output(['brew', '--prefix', 'llvm@20']).decode().strip()):
+    raise FileNotFoundError('LLVM not found, you can install it with `brew install llvm@20`')
   LLVM_PATH: str|None = os.path.join(brew_prefix, 'lib', 'libLLVM.dylib')
 else:
   LLVM_PATH = ctypes.util.find_library('LLVM')
   # use newer LLVM if possible
-  for ver in reversed(range(14, 19+1)):
+  for ver in reversed(range(14, 20+1)):
     if LLVM_PATH is not None: break
     LLVM_PATH = ctypes.util.find_library(f'LLVM-{ver}')
   if LLVM_PATH is None:

tinygrad/runtime/support/memory.py

@@ -0,0 +1,251 @@
+import collections, functools, dataclasses
+from typing import Any, ClassVar
+from tinygrad.helpers import round_up, getenv
+
+class TLSFAllocator:
+  """
+  The allocator is based on the Two-Level Segregated Fit (TLSF) algorithm. The allocator maintains 2 level of buckets:
+    * 1st level is determined by the most significant bit of the size.
+    * 2nd level splits the covered memory of 1st level into @lv2_cnt entries.
+
+  For each allocation request, the allocator searches for the smallest block that can fit the requested size.
+  For each deallocation request, the allocator merges the block with its neighbors if they are free.
+  """
+
+  def __init__(self, size:int, base:int=0, block_size:int=16, lv2_cnt:int=16):
+    self.size, self.base, self.block_size, self.l2_cnt = size, base, block_size, lv2_cnt.bit_length()
+    self.storage:list = [collections.defaultdict(list) for _ in range(size.bit_length() + 1)]
+    self.lv1_entries:list[int] = [0] * len(self.storage)
+
+    # self.blocks is more like a linked list, where each entry is a contiguous block.
+    self.blocks:dict[int, tuple[int, int|None, int|None, bool]] = {0: (size, None, None, True)} # size, next, prev, is_free
+    self._insert_block(0, size)
+
+  @functools.cache
+  def lv1(self, size): return size.bit_length()
+
+  @functools.cache
+  def lv2(self, size): return (size - (1 << (size.bit_length() - 1))) // (1 << max(0, size.bit_length() - self.l2_cnt))
+
+  def _insert_block(self, start:int, size:int, prev:int|None=None):
+    if prev is None: prev = self.blocks[start][2]
+    self.storage[self.lv1(size)][self.lv2(size)].append(start)
+    self.lv1_entries[self.lv1(size)] += 1
+    self.blocks[start] = (size, start + size, prev, True)
+    return self
+
+  def _remove_block(self, start:int, size:int, prev:int|None=None):
+    if prev is None: prev = self.blocks[start][2]
+    self.storage[self.lv1(size)][self.lv2(size)].remove(start)
+    self.lv1_entries[self.lv1(size)] -= 1
+    self.blocks[start] = (size, start + size, prev, False)
+    return self
+
+  def _split_block(self, start:int, size:int, new_size:int):
+    nxt = self.blocks[start][1]
+    assert self.blocks[start][3], "block must be free"
+    self._remove_block(start, size)._insert_block(start, new_size)._insert_block(start + new_size, size - new_size, prev=start)
+    if nxt in self.blocks: self.blocks[nxt] = (self.blocks[nxt][0], self.blocks[nxt][1], start + new_size, self.blocks[nxt][3])
+    return self
+
+  def _merge_right(self, start:int):
+    size, nxt, _, is_free = self.blocks[start]
+    assert is_free, "block must be free"
+
+    while is_free and nxt in self.blocks:
+      if (blk:=self.blocks[nxt])[3] is False: break
+      self._remove_block(start, size)._remove_block(nxt, blk[0])._insert_block(start, size:=size + blk[0])
+      assert self.blocks[start][1] == blk[1]
+      _, nxt, _, _ = self.blocks.pop(nxt)
+
+    if nxt in self.blocks: self.blocks[nxt] = (self.blocks[nxt][0], self.blocks[nxt][1], start, self.blocks[nxt][3])
+
+  def _merge_block(self, start:int):
+    # Go left while blocks are free. Then merge all them right.
+    while (x:=self.blocks[start][2]) is not None and self.blocks[x][3] is True: start = x
+    self._merge_right(start)
+
+  def alloc(self, req_size:int, align:int=1) -> int:
+    req_size = max(self.block_size, req_size) # at least block size.
+    size = max(self.block_size, req_size + align - 1)
+
+    # Round up the allocation size to the next bucket, so any entry there can fit the requested size.
+    size = round_up(size, (1 << size.bit_length() - self.l2_cnt))
+
+    # Search for the smallest block that can fit the requested size. Start with the it's bucket and go up until any block is found.
+    for l1 in range(self.lv1(size), len(self.storage)):
+      if self.lv1_entries[l1] == 0: continue
+      for l2 in range(self.lv2(size) if l1 == size.bit_length() else 0, (1 << self.l2_cnt)):
+        if len(self.storage[l1][l2]) > 0:
+          nsize = self.blocks[self.storage[l1][l2][0]][0]
+          assert nsize >= size, "block must be larger"
+
+          # Block start address.
+          start = self.storage[l1][l2][0]
+
+          # If request contains alignment, split the block into two parts.
+          if (new_start:=round_up(start, align)) != start:
+            self._split_block(start, nsize, new_start - start)
+            start, nsize = new_start, self.blocks[new_start][0]
+
+          # If the block is larger than the requested size, split it into two parts.
+          if nsize > req_size: self._split_block(start, nsize, req_size)
+          self._remove_block(start, req_size) # Mark the block as allocated.
+          return start + self.base
+    raise MemoryError(f"Can't allocate {req_size} bytes")
+
+  def free(self, start:int):
+    self._insert_block(start - self.base, self.blocks[start - self.base][0])._merge_block(start - self.base)
+
+# Memory Managment
+
+@dataclasses.dataclass(frozen=True)
+class VirtMapping: va_addr:int; size:int; paddrs:list[tuple[int, int]]; uncached:bool=False; system:bool=False; snooped:bool=False # noqa: E702
+
+class PageTableTraverseContext:
+  def __init__(self, dev, pt, vaddr, create_pts=False, free_pts=False, boot=False):
+    self.dev, self.vaddr, self.create_pts, self.free_pts, self.boot = dev, vaddr - dev.mm.va_base, create_pts, free_pts, boot
+    self.pt_stack:list[tuple[Any, int, int]] = [(pt, self._pt_pte_idx(pt, self.vaddr), self._pt_pte_size(pt))]
+
+  def _pt_pte_cnt(self, lv): return self.dev.mm.pte_cnt[lv]
+  def _pt_pte_size(self, pt): return self.dev.mm.pte_covers[pt.lv]
+  def _pt_pte_idx(self, pt, va): return (va // self._pt_pte_size(pt)) % self._pt_pte_cnt(pt.lv)
+
+  def level_down(self):
+    pt, pte_idx, _ = self.pt_stack[-1]
+
+    if not pt.valid(pte_idx):
+      assert self.create_pts, "Not allowed to create new page table"
+      pt.set_entry(pte_idx, self.dev.mm.palloc(0x1000, zero=True, boot=self.boot), table=True, valid=True)
+
+    assert not pt.is_huge_page(pte_idx), f"Must be table pt={pt.paddr:#x}, {pt.lv=} {pte_idx=} {pt.read_fields(pte_idx)}"
+    child_page_table = self.dev.mm.pt_t(self.dev, pt.address(pte_idx), lv=pt.lv+1)
+
+    self.pt_stack.append((child_page_table, self._pt_pte_idx(child_page_table, self.vaddr), self._pt_pte_size(child_page_table)))
+    return self.pt_stack[-1]
+
+  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)
+      parent_pt, parent_pte_idx, _ = self.pt_stack[-2]
+      parent_pt.set_entry(parent_pte_idx, 0x0, valid=False)
+      return True
+    return False
+
+  def level_up(self):
+    while self._try_free_pt() or self.pt_stack[-1][1] == self._pt_pte_cnt(self.pt_stack[-1][0].lv):
+      pt, pt_cnt, _ = self.pt_stack.pop()
+      if pt_cnt == self._pt_pte_cnt(pt.lv): self.pt_stack[-1] = (self.pt_stack[-1][0], self.pt_stack[-1][1] + 1, self.pt_stack[-1][2])
+
+  def next(self, size:int, paddr:int|None=None, off:int=0):
+    while size > 0:
+      pt, pte_idx, pte_covers = self.pt_stack[-1]
+      if self.create_pts:
+        assert paddr is not None, "paddr must be provided when allocating new page tables"
+        while pte_covers > size or not pt.supports_huge_page(paddr+off) or self.vaddr&(pte_covers-1) != 0: pt, pte_idx, pte_covers = self.level_down()
+      else:
+        while not pt.is_huge_page(pte_idx): pt, pte_idx, pte_covers = self.level_down()
+
+      entries = min(size // pte_covers, self._pt_pte_cnt(pt.lv) - pte_idx)
+      assert entries > 0, f"Invalid entries {size=:#x}, {pte_covers=:#x}"
+      yield off, pt, pte_idx, entries, pte_covers
+
+      size, off, self.vaddr = size - entries * pte_covers, off + entries * pte_covers, self.vaddr + entries * pte_covers
+      self.pt_stack[-1] = (pt, pte_idx + entries, pte_covers)
+      self.level_up()
+
+class MemoryManager:
+  va_allocator: ClassVar[TLSFAllocator|None] = None
+
+  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_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, self.level_cnt, self.va_bits = pt_t, palloc_ranges, len(va_shifts), va_bits
+
+    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
+    self.root_page_table = pt_t(self.dev, self.palloc(0x1000, zero=not self.dev.smi_dev, boot=True), lv=first_lv)
+
+  def _frag_size(self, va, sz, must_cover=True):
+    """
+    Calculate the tlb fragment size for a given virtual address and size.
+    If must_cover is True, the fragment size must cover the size, otherwise the biggest fragment size that fits the size is returned.
+    Fragment 0 is 4KB, 1 is 8KB and so on.
+    """
+    va_pwr2_div, sz_pwr2_div, sz_pwr2_max = va & -(va) if va > 0 else (1 << 63), sz & -(sz), (1 << (sz.bit_length() - 1))
+    return (min(va_pwr2_div, sz_pwr2_div) if must_cover else min(va_pwr2_div, sz_pwr2_max)).bit_length() - 1 - 12
+
+  def page_tables(self, vaddr:int, size:int):
+    ctx = PageTableTraverseContext(self.dev, self.root_page_table, vaddr, create_pts=True)
+    for _ in ctx.next(size, paddr=0): return [pt for pt, _, _ in ctx.pt_stack]
+
+  def map_range(self, vaddr:int, size:int, paddrs:list[tuple[int, int]], uncached=False, system=False, snooped=False, boot=False) -> VirtMapping:
+    if getenv("MM_DEBUG", 0): print(f"mm {self.dev.devfmt}: mapping {vaddr=:#x} ({size=:#x})")
+
+    assert size == sum(p[1] for p in paddrs), f"Size mismatch {size=} {sum(p[1] for p in paddrs)=}"
+
+    ctx = PageTableTraverseContext(self.dev, self.root_page_table, vaddr, create_pts=True, boot=boot)
+    for paddr, psize in paddrs:
+      for off, pt, pte_idx, pte_cnt, pte_covers in ctx.next(psize, paddr=paddr):
+        for pte_off in range(pte_cnt):
+          assert not pt.valid(pte_idx + pte_off), f"PTE already mapped: {pt.entry(pte_idx + pte_off):#x}"
+          pt.set_entry(pte_idx + pte_off, paddr + off + pte_off * pte_covers, uncached=uncached, system=system, snooped=snooped,
+                       frag=self._frag_size(ctx.vaddr+off, pte_cnt * pte_covers), valid=True)
+
+    self.on_range_mapped()
+    return VirtMapping(vaddr, size, paddrs, uncached=uncached, system=system, snooped=snooped)
+
+  def unmap_range(self, vaddr:int, size:int):
+    if getenv("MM_DEBUG", 0): print(f"mm {self.dev.devfmt}: unmapping {vaddr=:#x} ({size=:#x})")
+
+    ctx = PageTableTraverseContext(self.dev, self.root_page_table, vaddr, free_pts=True)
+    for off, pt, pte_idx, pte_cnt, pte_covers in ctx.next(size):
+      for pte_id in range(pte_idx, pte_idx + pte_cnt):
+        assert pt.valid(pte_id), f"PTE not mapped: {pt.entry(pte_id):#x}"
+        pt.set_entry(pte_id, paddr=0x0, valid=False)
+
+  def on_range_mapped(self): pass
+
+  @classmethod
+  def alloc_vaddr(cls, size:int, align=0x1000) -> int:
+    assert cls.va_allocator is not None, "must be set it"
+    return cls.va_allocator.alloc(size, max((1 << (size.bit_length() - 1)), align))
+
+  def valloc(self, size:int, align=0x1000, uncached=False, contiguous=False) -> VirtMapping:
+    # Alloc physical memory and map it to the virtual address
+    va = self.alloc_vaddr(size:=round_up(size, 0x1000), align)
+
+    if contiguous: 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.
+      nxt_range, rem_size, paddrs = 0, size, []
+      while rem_size > 0:
+        while self.palloc_ranges[nxt_range][0] > rem_size: nxt_range += 1
+
+        try: paddrs += [(self.palloc(try_sz:=self.palloc_ranges[nxt_range][0], self.palloc_ranges[nxt_range][1], zero=False), try_sz)]
+        except MemoryError:
+          # Move to a smaller size and try again.
+          nxt_range += 1
+          if nxt_range == len(self.palloc_ranges):
+            for paddr, _ in paddrs: self.pa_allocator.free(paddr)
+            raise MemoryError(f"Failed to allocate memory. (total allocation size={size:#x}, current try={self.palloc_ranges[nxt_range-1]})")
+          continue
+        rem_size -= self.palloc_ranges[nxt_range][0]
+
+    return self.map_range(va, size, paddrs, uncached=uncached)
+
+  def vfree(self, vm:VirtMapping):
+    assert self.va_allocator is not None, "must be set it"
+    self.unmap_range(vm.va_addr, vm.size)
+    self.va_allocator.free(vm.va_addr)
+    for paddr, _ in vm.paddrs: self.pa_allocator.free(paddr)
+
+  def palloc(self, size:int, align:int=0x1000, zero=True, boot=False) -> int:
+    assert self.dev.is_booting == boot, "During booting, only boot memory can be allocated"
+    paddr = (self.boot_allocator if boot else self.pa_allocator).alloc(round_up(size, 0x1000), align)
+    if zero: self.dev.vram[paddr:paddr+size] = bytes(size)
+    return paddr
+
+  def pfree(self, paddr:int): self.pa_allocator.free(paddr)