From 1bf0a7a2d15938b365d9fb71565a55e82a02e21f Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Thu, 28 Mar 2024 20:26:38 -0700
Subject: [PATCH] move assign logic into lazy.py (#3984)

* move assign logic into lazy.py

* don't check the buffer
---
 docs/abstractions2.py      | 2 ++
 tinygrad/engine/realize.py | 9 ++-------
 tinygrad/lazy.py           | 6 ++++--
 tinygrad/tensor.py         | 2 ++
 4 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/docs/abstractions2.py b/docs/abstractions2.py
index 2bfbbaf3c5..7888bdcb25 100644
--- a/docs/abstractions2.py
+++ b/docs/abstractions2.py
@@ -81,6 +81,8 @@ a = LazyBuffer.loadop(LoadOps.EMPTY, (1,), dtypes.int32, DEVICE)
 b = LazyBuffer.loadop(LoadOps.EMPTY, (1,), dtypes.int32, DEVICE)
 a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
 b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
+del a.srcs
+del b.srcs
 
 # describe the computation
 out = a.e(BinaryOps.ADD, b)
diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py
index 60ebfcacaa..0b1fa1d0c8 100644
--- a/tinygrad/engine/realize.py
+++ b/tinygrad/engine/realize.py
@@ -43,13 +43,8 @@ def run_schedule(schedule:List[ScheduleItem]):
 
     for out in si.outputs:
       # we don't have an output buffer, we have to create it, and create to max size if it has symbolic shape
-      if out.size > 0:
-        if out.op is LoadOps.ASSIGN and out.srcs[1].base.realized is not None:
-          # if the buffer isn't realized, it might be a const or something. this is fine
-          out.buffer = out.srcs[1].base.buffer
-        else:
-          if not dont_allocate: out.buffer.allocate()
-        del out.srcs
+      if out.size > 0 and not dont_allocate and out.op is not LoadOps.ASSIGN: out.buffer.allocate()
+      del out.srcs
 
     # run the function (put it in JIT)
     real_buffers = [x.buffer for x in si.outputs+si.inputs if x.size != 0]
diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py
index b743725875..31c2655be2 100644
--- a/tinygrad/lazy.py
+++ b/tinygrad/lazy.py
@@ -33,7 +33,8 @@ class LazyBuffer:
     if base is None:
       # properties on base
       self.op, self.arg, self.srcs = op, arg, srcs  # this is a LazyOp, except the src is LazyBuffers and not LazyOps
-      self.buffer: Buffer = Buffer(device, self.size, dtype)
+      assert self.op is not LoadOps.ASSIGN or srcs[1].base.realized is not None, "assign target must be realized"
+      self.buffer: Buffer = srcs[1].base.buffer if self.op is LoadOps.ASSIGN else Buffer(device, self.size, dtype)
       self.contiguous_child: Optional[Tuple[ReferenceType[LazyBuffer], ShapeTracker]] = None
       self.forced_realize = False
     else:
@@ -46,7 +47,8 @@ class LazyBuffer:
 
   @property
   def realized(self) -> Optional[Buffer]:
-    return self.buffer if self._base is None and hasattr(self.buffer, "_buf") else None
+    # NOTE: we check for a lack of srcs instead of an allocated buffer to make unrealized assigns return None here
+    return self.buffer if self._base is None and not hasattr(self, 'srcs') else None
 
   # NOTE: this has to be a function to prevent self reference
   @property
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index a7f6f41c5d..8c1e71f95f 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -46,7 +46,9 @@ def _loadop(op, shape:Tuple[sint,...], dtype:DType, device:Union[str, Tuple[str,
 
 def _fromcpu(x: np.ndarray) -> LazyBuffer:
   ret = LazyBuffer.loadop(LoadOps.EMPTY, x.shape, dtypes.from_np(x.dtype), "EXT")
+  # fake realize
   ret.buffer.allocate((memoryview(bytearray()), None) if x.size == 0 else (flat_mv(np.require(x, requirements='C').data), x))
+  del ret.srcs
   return ret
 
 def _get_winograd_matcols(mat, dims:int, shp:Tuple[sint, ...], device:Union[str, Tuple[str, ...]]) -> List[List[Tensor]]: