github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 15:38:29 -05:00
Code Issues Packages Projects Releases Wiki Activity

nv do not map regions twice (#5521)

Browse Source
This commit is contained in:
nimlgen
2024-07-17 11:20:02 +03:00
committed by GitHub
parent 2d53abb04a
commit 661da32aff
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tinygrad/runtime/ops_nv.py
View File

@@ -451,13 +451,13 @@ class NVDevice(HCQCompatCompiled):
# NOTE: va_addr is set to make rawbufs compatable with AMD.
return uvm.map_external_allocation(self.fd_uvm, base=va_base, length=size, rmCtrlFd=self.fd_ctl, hClient=self.root, hMemory=mem_handle,
gpuAttributesCount=1, perGpuAttributes=gpu_attrs, va_addr=va_base, size=size)
gpuAttributesCount=1, perGpuAttributes=gpu_attrs, va_addr=va_base, size=size, mapped_gpu_ids=[self.gpu_uuid])
def _gpu_map(self, mem):
mem = mem._base if hasattr(mem, '_base') else mem
if self.gpu_uuid in getattr(mem, "mapped_gpu_ids", []): return
mem.__setattr__("mapped_gpu_ids", getattr(mem, "mapped_gpu_ids", []) + [self.gpu_uuid])
return self._gpu_uvm_map(mem.va_addr, mem.size, mem.hMemory, create_range=False)
if self.gpu_uuid in mem.mapped_gpu_ids: return
mem.mapped_gpu_ids.append(self.gpu_uuid)
self._gpu_uvm_map(mem.va_addr, mem.size, mem.hMemory, create_range=False)
def _alloc_gpu_vaddr(self, size, alignment=(4 << 10)):
NVDevice.uvm_vaddr = (res_va:=round_up(NVDevice.uvm_vaddr, alignment)) + size