From 2c324d0685323fca8a2d78e5ec9ff5ea1df9b0c3 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Fri, 9 Jun 2023 21:28:06 -0700
Subject: [PATCH] fix metal uaf (#964)

---
 test/external/external_metal_uaf.py | 17 +++++++++++++++++
 tinygrad/runtime/lib.py             |  2 +-
 2 files changed, 18 insertions(+), 1 deletion(-)
 create mode 100644 test/external/external_metal_uaf.py

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