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

refactor to trace_num property in buffers (#12638)

Browse Source
This commit is contained in:
qazal
2025-10-13 18:06:55 +08:00
committed by GitHub
parent cd6aeebfee
commit 066d25f5fb
2 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/device.py
+6 -4
View File
@@ -137,16 +137,14 @@ class Buffer:
else:
self._buf = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options)
if not self.device.startswith("DISK"): GlobalCounters.mem_used += self.nbytes
if PROFILE:
self._prof_num = num = len(Buffer.profile_events)
Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", num, {"dtype":self.dtype, "sz":self.size}))
if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size}))
return self
def deallocate(self):
assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"
if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
if self._base is None and (self.options is None or self.options.external_ptr is None):
if GlobalCounters is not None and not self.device.startswith("DISK"): GlobalCounters.mem_used -= self.nbytes
if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self._prof_num))
if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num))
self.allocator.free(self._buf, self.nbytes, self.options)
elif self._base is not None: self._base.allocated_views -= 1
del self._buf
@@ -160,6 +158,10 @@ class Buffer:
self.copyout(memoryview(buf))
return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount)
@property
def trace_num(self) -> int:
if not hasattr(self, '_trace_num'): self._trace_num = len(Buffer.profile_events)
return self._trace_num
@property
def nbytes(self): return self.size*self.dtype.itemsize
def __del__(self): (not hasattr(self, '_buf')) or self.deallocate()
def __repr__(self):
tinygrad/engine/realize.py
+3 -1
View File
@@ -165,7 +165,9 @@ class ExecItem:
def run(self, _var_vals:dict[str, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None:
var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars)
bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs]
if PROFILE: cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, {"metadata":self.metadata, "var_vals":var_vals}))
if PROFILE:
payload = {"metadata":self.metadata, "var_vals":var_vals, "bufs":[b.trace_num for b in bufs]}
cpu_events.append(ProfilePointEvent(self.prg.device, "exec", self.prg.display_name, payload))
et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2)
if do_update_stats:
GlobalCounters.kernel_count += 1