From eb7df9213699dcd0ab560d27ac988809b77e337d Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Mon, 6 Jan 2025 10:37:20 +0200
Subject: [PATCH] dedup COPY UOp [pr] (#8506)

---
 test/test_schedule.py | 12 ++++++------
 tinygrad/ops.py       |  1 +
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/test/test_schedule.py b/test/test_schedule.py
index a9e5cf56fd..69c9e9ca9e 100644
--- a/test/test_schedule.py
+++ b/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,))
diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index dbdc100e71..70d10de95b 100644
--- a/tinygrad/ops.py
+++ b/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: