move assign logic into lazy.py (#3984)

* move assign logic into lazy.py

* don't check the buffer

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)

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]

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

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]]: