copy movement op tests construct copies (#17792)

* copy movement op tests construct copies

* hop 3 gpus

* assert the amount of bytes moved

* work

* arange 4, everyone copies 16B first

test/backend/test_schedule.py

@@ -379,27 +379,30 @@ class TestCopyFolding(unittest.TestCase):
     check_schedule(a.clone(), 1, filter_sink=False)
 
   def test_shrink_copy(self):
-    a = Tensor.arange(4)
-    view = a.shrink(((0, 2),))
-    b = view.clone()
-    run_linear(*check_schedule(b, 1, filter_sink=False))
-    self.assertEqual(b.uop.base.buffer.size, 2)
-    self.assertEqual(b.uop.numel(), 2)
-    self.assertListEqual(b.tolist(), [0, 1])
+    a = Tensor.arange(4).clone("CPU:1").realize()
+    b = a.to("CPU:2").shrink(((1, 3),)).to("CPU:3")
+    GlobalCounters.reset()
+    run_linear(*check_schedule(b, 3, filter_sink=False))
+    # extra E kernel, copy exactly 4 bytes
+    self.assertEqual(GlobalCounters.global_mem, 4*4 + 2*4*2 + 2*4)
+    self.assertListEqual(b.tolist(), [1, 2])
 
   def test_expanded_copy(self):
-    a = Tensor.arange(2)
-    view = a.reshape(2, 1).expand(2, 2)
-    b = view.clone()
-    run_linear(*check_schedule(b, 1, filter_sink=False))
-    self.assertEqual(b.uop.base.buffer.size, 4)
-    self.assertEqual(b.uop.numel(), 4)
-    self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])
+    a = Tensor.arange(4).clone("CPU:1").realize()
+    b = a.to("CPU:2").reshape(4, 1).expand(4, 2).to("CPU:3")
+    GlobalCounters.reset()
+    run_linear(*check_schedule(b, 3, filter_sink=False))
+    # TODO: expands before copy
+    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 8*4) + 8*4)
+    self.assertListEqual(b.tolist(), [[0, 0], [1, 1], [2, 2], [3, 3]])
 
   def test_permuted_copy(self):
-    a = Tensor.arange(4)
-    b = a.reshape(2, 2).permute(1, 0)
-    b.realize()
+    a = Tensor.arange(4).clone("CPU:1").realize()
+    b = a.to("CPU:2").reshape(2, 2).permute(1, 0).to("CPU:3")
+    GlobalCounters.reset()
+    run_linear(*check_schedule(b, 3, filter_sink=False))
+    # permutes before copy
+    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 4*4) + 4*4)
     self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
 
   def test_permute_on_disk(self):