nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

simpler hcq buffer (#17148)

Browse Source 
* simpleer hcq buffer

* x
This commit is contained in:
nimlgen
2026-07-23 13:28:51 +03:00
committed by GitHub
parent c5b2b9242d
commit 62273d50fc
6 changed files with 24 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/hcq2/hcq2.py
+15 -21
View File
@@ -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")
tinygrad/device.py
+4 -2
View File
@@ -191,9 +191,11 @@ class Buffer:
def __repr__(self):
return f"<buf real:{self.is_allocated()} device:{self.device} size:{self.size} dtype:{self.dtype}" + \
(f" offset:{self.offset}" if self._base is not None else "") + (f" {self.options=}" if self.options is not None else "") + ">"
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
tinygrad/engine/realize.py
+1 -1
View File
@@ -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
tinygrad/runtime/ops_cpu.py
+1 -3
View File
@@ -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)
tinygrad/runtime/ops_metal.py
+2 -3
View File
@@ -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)
tinygrad/runtime/ops_qcom.py
+1 -3
View File
@@ -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)