tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 17:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

am: fix mmap failures (#10054)

Browse Source
This commit is contained in:
nimlgen
2025-04-26 14:21:28 +03:00
committed by GitHub
parent e1d2b64e92
commit e055b9422f
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/runtime/ops_amd.py
+6 -3
View File
@@ -647,7 +647,7 @@ class KFDIface:
raise RuntimeError("\n".join(report))
@dataclass
class AMAllocationMeta: owner:AMDDevice; mapped_devs:list[AMDDevice]; mapping:AMMapping # noqa: E702
class AMAllocationMeta: owner:AMDDevice; mapped_devs:list[AMDDevice]; mapping:AMMapping; has_cpu_mapping:bool # noqa: E702
class PCIIface:
supported_devs:list[int] = [0x744c, 0x7480, 0x7550]
@@ -737,26 +737,29 @@ class PCIIface:
def _map_pci_range(self, bar, off=0, addr=0, size=None, fmt='B'):
fd, sz = self.bar_fds[bar], size or (self.bar_info[bar][1] - self.bar_info[bar][0] + 1)
libc.madvise(loc:=fd.mmap(addr, sz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if addr else 0), off), sz, libc.MADV_DONTFORK)
assert loc != 0xffffffffffffffff, f"Failed to mmap {size} bytes at {hex(addr)}"
return MMIOInterface(loc, sz, fmt=fmt)
def alloc(self, size:int, host=False, uncached=False, cpu_access=False):
if host or (not getenv("AMD_ALLOC_QUEUE_DEV_MEM", 1) and uncached and cpu_access): # host or gtt-like memory.
vaddr = self.adev.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
va = FileIOInterface.anon_mmap(vaddr, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS | MAP_LOCKED | MAP_FIXED, 0)
assert va != 0xffffffffffffffff, f"Failed to mmap {size} bytes at {hex(vaddr)}"
# Read pagemap to get the physical address of each page. The pages are locked.
self.pagemap.seek(va // mmap.PAGESIZE * 8)
paddrs = [((x & ((1<<55) - 1)) * mmap.PAGESIZE, mmap.PAGESIZE) for x in array.array('Q', self.pagemap.read(size//mmap.PAGESIZE*8, binary=True))]
am_mapping = self.adev.mm.map_range(vaddr, size, paddrs, system=True, snooped=True, uncached=True)
return HCQBuffer(vaddr, size, meta=AMAllocationMeta(self.dev, [self.dev], am_mapping))
return HCQBuffer(vaddr, size, meta=AMAllocationMeta(self.dev, [self.dev], am_mapping, has_cpu_mapping=cpu_access))
am_mapping = self.adev.mm.valloc(size:=round_up(size, 4 << 10), uncached=uncached, contigous=cpu_access)
if cpu_access: self._map_pci_range(bar=0, off=am_mapping.paddrs[0][0], addr=am_mapping.va_addr, size=am_mapping.size)
return HCQBuffer(am_mapping.va_addr, size, meta=AMAllocationMeta(self.dev, [self.dev], am_mapping))
return HCQBuffer(am_mapping.va_addr, size, meta=AMAllocationMeta(self.dev, [self.dev], am_mapping, has_cpu_mapping=cpu_access))
def free(self, mem):
for dev in mem.meta.mapped_devs[1:]: dev.dev_iface.adev.mm.unmap_range(mem.va_addr, mem.size)
if not mem.meta.mapping.system: self.adev.mm.vfree(mem.meta.mapping)
if mem.meta.owner == self.dev and mem.meta.has_cpu_mapping: FileIOInterface.munmap(mem.va_addr, mem.size)
def map(self, mem):
# Check if the memory is already mapped on this device