fix metal uaf (tinygrad#964)

awstone · Jun 10, 2023 · 2c324d0 · 2c324d0
1 parent c0e558b
commit 2c324d0
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/test/external/external_metal_uaf.py 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
diff --git a/tinygrad/runtime/lib.py b/tinygrad/runtime/lib.py
@@ -31,7 +31,7 @@ def fromCPU(cls, x:np.ndarray, **kwargs):
 
 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