Method valloc

tinygrad/runtime/support/memory.py:239–264

(self, size:int, align=0x1000, uncached=False, contiguous=False)

    return va

  def valloc(self, size:int, align=0x1000, uncached=False, contiguous=False) -> VirtMapping:
    if not getenv("GMMU", 1):
      paddr = self.palloc(size:=round_up(size, 0x1000), align, zero=False)
      return VirtMapping(self.identity_va(uncached) + paddr, size, [(paddr, size)], aspace=AddrSpace.PHYS, uncached=uncached)

    # 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.pfree(paddr)
            raise MemoryError(f"Failed to allocate memory (OOM). Request size={size:#x} ({self.palloc_ranges[nxt_range-1]})")
          continue
        rem_size -= self.palloc_ranges[nxt_range][0]

    return self.map_range(va, size, paddrs, aspace=AddrSpace.PHYS, uncached=uncached)

  def vfree(self, vm:VirtMapping):
    if not getenv("GMMU", 1): return self.pfree(vm.paddrs[0][0])