diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index fbdef9f444..6d69671d81 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -297,6 +297,8 @@ class AMDAllocator(HCQAllocator['AMDDevice']): def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf) + def _do_unmap(self, buf:HCQ2Buffer): self.dev.iface.unmap(buf) + @dataclass class AMDQueueDesc: ring: Buffer; read_ptr: Buffer; write_ptr: Buffer; doorbell: Buffer; put_value: Buffer # noqa: E702 @@ -388,15 +390,24 @@ class KFDIface: return hcqbuf def free(self, mem): + self._unmap(mem) + if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size) + kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle) + + def unmap(self, mem): + self._unmap(mem) + if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle) + + def _unmap(self, mem): gpus = (ctypes.c_int32 * 1)(self.gpu_id) stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1) assert stm.n_success == 1 - if mem.owner == self.dev: - if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size) - kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle) def map(self, mem): - if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr) + if mem.owner is not None and mem.owner._is_cpu(): + mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr) + mapped._owns_kfd_handle = True + return mapped c_gpus = (ctypes.c_int32 * 1)(self.gpu_id) stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1) @@ -468,6 +479,7 @@ class PCIIface(PCIIfaceBase): def require_profile_mode(self): return True def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics. + def unmap(self, mem): self.free(mem) def _compute_props(self): self.ip_versions = self.dev_impl.ip_ver diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 767291e4f8..2464928ca5 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -623,6 +623,8 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]): if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented") return self._do_map(buf) + def _do_unmap(self, mb): self.dev.iface.free(mb) + @suppress_finalizing def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None): if options is not None and options.external_ptr is not None: return @@ -631,6 +633,6 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]): def _unmap(self, mb): self.dev.synchronize() - self.dev.iface.free(mb) + self._do_unmap(mb) def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)