diff --git a/test/external/external_metal_uaf.py b/test/external/external_metal_uaf.py new file mode 100644 index 0000000000..46b89460e8 --- /dev/null +++ b/test/external/external_metal_uaf.py @@ -0,0 +1,17 @@ +import weakref +import numpy as np +from tinygrad.tensor import Tensor, Device +Device.DEFAULT = "METAL" + +if __name__ == "__main__": + t = Tensor.zeros(3).realize() + wt = weakref.ref(t.lazydata.realized) + n = t.numpy() + t += 1 + n2 = t.numpy() + print(wt) + del t + print(wt) + print(n, n.base, n.base.base) + print(n2, n2.base, n2.base.base) + assert wt() is not None \ No newline at end of file diff --git a/tinygrad/runtime/lib.py b/tinygrad/runtime/lib.py index 313e13b0a4..cbc4e73adc 100644 --- a/tinygrad/runtime/lib.py +++ b/tinygrad/runtime/lib.py @@ -31,7 +31,7 @@ class RawBufferCopyIn(RawBuffer): class RawBufferMapped(RawBufferCopyIn): def _buffer(self) -> memoryview: raise NotImplementedError("must be implemented") - def toCPU(self) -> np.ndarray: return np.frombuffer(self._buffer(), dtype=self.dtype.np) + def toCPU(self) -> np.ndarray: return np.frombuffer(self._buffer(), dtype=np.dtype(self.dtype.np, metadata={"backing": self})) # type: ignore def _copyin(self, x:np.ndarray) -> None: np.copyto(self.toCPU(), x.reshape(-1)) # this one is simple enough that i moved it out of the runtimes