diff --git a/extra/hcq2/hcq2.py b/extra/hcq2/hcq2.py index 58bddb7b0a..64ea59e7e1 100644 --- a/extra/hcq2/hcq2.py +++ b/extra/hcq2/hcq2.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import cast, Callable, TypeVar, Generic, Any import struct, functools, time, collections, itertools from dataclasses import replace, dataclass -from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE +from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp @@ -410,13 +410,14 @@ def push_stack(op, s): return UOp(Ops.STACK, op.dtype.scalar(), tuple(op.replace(dtype=op.dtype.scalar(), src=tuple(x if y is s else y for y in op.src)) for x in s.src)) def fold_binary(buf:UOp, blob:UOp) -> UOp: - for b in (m.bufs if isinstance(m:=buf.buffer, MultiBuffer) else (m,)): b.ensure_allocated()._buf.cpu_view().view(fmt='B')[:len(blob.arg)] = blob.arg + for b in (m.bufs if isinstance(m:=buf.buffer, MultiBuffer) else (m,)): + b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[:len(blob.arg)] = blob.arg return UOp(Ops.NOOP) def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp: for b, v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)): data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.arg)) - b.ensure_allocated()._buf.cpu_view().view(offset=off.arg * buf.dtype.itemsize, size=len(data), fmt='B')[:] = data + b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(byte_off:=off.arg*buf.dtype.itemsize):byte_off+len(data)] = data return UOp(Ops.NOOP) def resolve_getaddr(buf:UOp, g:UOp) -> UOp: @@ -492,19 +493,19 @@ class HCQ2Compiled(Compiled): @functools.cache def timeline_signal(self, queue:str|None=None, init_value:int=0) -> Buffer: buf = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True) - buf._buf.cpu_view().mv.cast('Q')[0] = init_value + buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value return buf @functools.cache def timeline_value(self, queue:str|None=None, init_value:int=1) -> Buffer: buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True) - buf.as_memoryview(force_zero_copy=True).cast('Q')[0] = init_value + buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value return buf def synchronize(self, timeout:int|None=None): if not hasattr(self, 'iface'): return - sig = self.timeline_signal()._buf.cpu_view().mv.cast('Q') - tl = self.timeline_value().as_memoryview(force_zero_copy=True).cast('Q') + sig = self.timeline_signal().as_memoryview(force_zero_copy=True, no_sync=True).cast('Q') + tl = self.timeline_value().as_memoryview(force_zero_copy=True, no_sync=True).cast('Q') st = time.perf_counter() while sig[0] < tl[0] - 1: if time.perf_counter() - st > (timeout or 3000) / 1000: self.on_device_hang() @@ -532,25 +533,18 @@ class HCQ2Compiled(Compiled): # if the device has an interface, call device_fini to clean up resources if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini() +@dataclass class HCQ2Buffer: - def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQ2Buffer|None=None, view:MMIOInterface|None=None, owner:HCQ2Compiled|None=None): - self.va_addr, self.size, self.meta, self._base, self.view, self.owner = va_addr, size, meta, _base, view, owner + va_addr:sint + meta:Any=None + view:MMIOInterface|None=None - def offset(self, offset:int=0, size:int|None=None) -> HCQ2Buffer: - return HCQ2Buffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, meta=self.meta, - _base=self._base or self, view=(self.view.view(offset=offset, size=size) if self.view is not None else None)) - - def cpu_view(self) -> MMIOInterface: - assert self.view is not None, "buffer has no cpu_view" - return self.view - - @property - def base(self) -> HCQ2Buffer: return self._base or self + def offset(self, offset:int, size:int) -> HCQ2Buffer: + return HCQ2Buffer(self.va_addr+offset, meta=self.meta, view=(self.view.view(offset=offset, size=size) if self.view is not None else None)) class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]): def _as_buffer(self, buf:HCQ2Buffer) -> memoryview: - self.dev.synchronize() - return buf.cpu_view().mv + return unwrap(buf.view).mv def _map(self, buf:HCQ2Buffer) -> HCQ2Buffer: if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented") diff --git a/tinygrad/device.py b/tinygrad/device.py index 738c8b0eca..a8dc8562cc 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -191,9 +191,11 @@ class Buffer: def __repr__(self): return f"" - def as_memoryview(self, allow_zero_copy=False, force_zero_copy=False) -> memoryview: + def as_memoryview(self, allow_zero_copy=False, force_zero_copy=False, no_sync=False) -> memoryview: # zero copy with as_memoryview (disabled by default due to use after free) - if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf) + if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer'): + if not no_sync: self.allocator.dev.synchronize() + return self.allocator._as_buffer(self._buf) assert not force_zero_copy, "force zero copy was passed, but copy is required" Buffer("PYTHON", self.size, self.dtype, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self) return mv diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 482133e3ce..d2e8e9dff1 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -168,7 +168,7 @@ def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \ and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk: dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes) - elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf) + elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf) else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True)) return None diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 7c28a72537..f9ec59932c 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -137,9 +137,7 @@ class CPUAllocator(HCQAllocator): elif WIN: addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE)) else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE)) return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev) - def _as_buffer(self, src) -> memoryview: - self.dev.synchronize() - return to_mv(src.va_addr, src.size) + def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size) def _do_map(self, buf:HCQBuffer): if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu") return HCQBuffer(buf.view.addr, buf.size, view=buf.view, owner=buf.owner) diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index 1e14bce6a2..4897c77965 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -183,10 +183,9 @@ class MetalAllocator(LRUAllocator[MetalDevice]): # There is no real metal multidevice support for now, so transfer is used only for tests. src_dev.synchronize() def _cp_mv(self, dst, src, prof_desc): - with cpu_profile(prof_desc, f"{self.dev.device}:COPY"): dst[:] = src - def _as_buffer(self, src:MetalBuffer) -> memoryview: self.dev.synchronize() - return to_mv(src.buf.contents(), src.size + src.offset)[src.offset:] + with cpu_profile(prof_desc, f"{self.dev.device}:COPY"): dst[:] = src + def _as_buffer(self, src:MetalBuffer) -> memoryview: return to_mv(src.buf.contents(), src.size + src.offset)[src.offset:] def _copyin(self, dest:MetalBuffer, src:memoryview): self._cp_mv(self._as_buffer(dest), src, "TINY -> METAL") def _copyout(self, dest:memoryview, src:MetalBuffer): self._cp_mv(dest, self._as_buffer(src), "METAL -> TINY") def _offset(self, buf:MetalBuffer, size:int, offset:int): return MetalBuffer(buf.buf, size, offset) diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index 8fed9fdb9b..21c6ded875 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -332,9 +332,7 @@ class QCOMAllocator(HCQAllocatorBase): def _copyin(self, dest:HCQBuffer, src:memoryview): self._do_copy(mv_address(src), dest.cpu_view().addr, src.nbytes, f"TINY -> {self.dev.device}") def _copyout(self, dest:memoryview, src:HCQBuffer): self._do_copy(src.cpu_view().addr, mv_address(dest), src.size, f"{self.dev.device} -> TINY") - def _as_buffer(self, src:HCQBuffer) -> memoryview: - self.dev.synchronize() - return to_mv(src.cpu_view().addr, src.size) + def _as_buffer(self, src:HCQBuffer) -> memoryview: return to_mv(src.cpu_view().addr, src.size) def _do_free(self, opaque, options:BufferSpec): self.dev._gpu_free(opaque)