dedup COPY UOp [pr] (#8506)

test/test_schedule.py

@@ -230,16 +230,16 @@ class TestSchedule(unittest.TestCase):
     # a and b share the same underlying device memory
     self.assertIs(a.lazydata.realized, b.lazydata.realized)
 
-  # EMPTY and COPY are assigned to unique device Buffers
-
-  def test_no_dedup_copy(self):
+  def test_copy_dedups(self):
     src = Tensor.ones(4).contiguous().realize()
     a = src.clone()
     b = src.clone()
-    sched = check_schedule([a, b], 2, filter_sink=False)
+    sched = check_schedule([a, b], 1, filter_sink=False)
     run_schedule(sched)
-    # a and b are assigned to different device Buffers
-    self.assertIsNot(a.lazydata.realized, b.lazydata.realized)
+    # a and b are assigned to the same device Buffer
+    self.assertIs(a.lazydata.realized, b.lazydata.realized)
+
+  # EMPTY is assigned to a unique device Buffer
 
   def test_no_dedup_empty(self):
     a = Tensor.empty((4,))

tinygrad/ops.py

@@ -444,6 +444,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if op is Ops.BIND:
       assert isinstance(arg, UOp) and arg.op is Ops.BIND and shape == (), f"trying to create BIND with {arg=} {shape=}"
       return UOp(Ops.VIEW, dtype, (UOp(Ops.DEVICE, arg=device), arg), ShapeTracker.from_shape(()))
+    if op is Ops.COPY: return UOp(op, dtype, src, arg)
     # otherwise it's a contiguous st
     return UOp(Ops.VIEW, dtype, (UOp.new_buffer(device, (st:=ShapeTracker.from_shape(shape)).size, dtype), UOp(op, dtype, src, arg)), st)
   def copy_to_device(self, device:str, force=False, clone:bool=False) -> UOp: