Draft: port sliced-assign fast path onto upstream 2026-09-28

test/runtime/test_setitem_schedule.py

@@ -122,5 +122,44 @@ def test_setitem_slice_assign_from_other_device(self):
     assert_kernel_count(3 if Device.DEFAULT in {"CL", "WEBGPU"} else 2)
     self.assertListEqual(a.tolist(), [1.0]*10 + [13.0, 14.0] + [1.0]*8)
 
+  def test_cross_device_slice_assign_single_copy(self):
+    a = Tensor.ones(20, device="CPU").contiguous().realize()
+    b = Tensor.arange(20, dtype=dtypes.float32).clone("CPU:1").realize()
+    GlobalCounters.reset()
+    a[10:12].assign(b[13:15].to(a.device)).realize()
+    assert_kernel_count(1)
+    self.assertEqual(GlobalCounters.global_mem, 2*4)
+    self.assertListEqual(a.tolist(), [1.0]*10 + [13.0, 14.0] + [1.0]*8)
+
+  def test_cross_device_slice_assign_composed_contiguous_views(self):
+    for width in (1, 3, 8):
+      for start, stop in ((0, 1), (1, 3), (3, 4), (0, 4)):
+        for layout in ("rows", "nested", "singleton_permute", "inverse_permute"):
+          with self.subTest(width=width, start=start, stop=stop, layout=layout):
+            size = (stop-start)*width
+            before = list(range(4*width))
+            values = list(range(100, 100+size+2))
+            dst = Tensor(before, dtype=dtypes.int32, device="CPU").realize()
+            src = Tensor(values, dtype=dtypes.int32, device="CPU:1").realize()
+            view = dst.reshape(4, width)[start:stop]
+            if layout == "nested": view = dst.reshape(2, 2, width).reshape(4, width)[start:][:(stop-start)]
+            if layout == "singleton_permute": view = view.reshape(stop-start, 1, width).permute(1, 0, 2)
+            if layout == "inverse_permute": view = view.permute(1, 0).permute(1, 0)
+            GlobalCounters.reset()
+            view.assign(src[1:1+size].reshape(view.shape).to(view.device)).realize()
+            assert_kernel_count(1)
+            expected = before[:]
+            expected[start*width:stop*width] = values[1:1+size]
+            self.assertListEqual(dst.tolist(), expected)
+            self.assertListEqual(src.tolist(), values)
+
+  def test_cross_device_slice_assign_noncontiguous_fallback(self):
+    a = Tensor.zeros(4, 4, device="CPU").contiguous().realize()
+    b = Tensor.arange(4, dtype=dtypes.float32).reshape(2, 2).clone("CPU:1").realize()
+    GlobalCounters.reset()
+    a[1:3, 1:3].assign(b.to(a.device)).realize()
+    assert_kernel_count(2)
+    self.assertListEqual(a.tolist(), [[0.0]*4, [0.0, 0.0, 1.0, 0.0], [0.0, 2.0, 3.0, 0.0], [0.0]*4])
+
 if __name__ == '__main__':
   unittest.main()

tinygrad/schedule/__init__.py

@@ -228,6 +228,9 @@ def collect_stores(ctx:CallifyCtx, u:UOp):
   # fold MOPS+BITCAST over BUFFER into SHRINK when movement ops collapse to contiguous range
   (UPat((Ops.COPY, Ops.STAGE), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), allow_any_len=True, name="c"), contiguous_mops_to_view),
   (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, name="src"), UPat()), name="c", allow_any_len=True), contiguous_mops_to_view),
+  # A cross-device copy can target the contiguous destination view directly, avoiding a separate assign kernel.
+  (UPat(Ops.STORE, src=(UPat(GroupOp.Movement, name="src"), UPat(Ops.COPY, name="cpy")), name="c"),
+   lambda ctx,c,src,cpy: contiguous_mops_to_view(ctx, c, src) if src.base.op is Ops.BUFFER and cpy.src[0].device != c.device else None),
 
   # Collect effects after their sources have been rewritten, without entering call bodies.
   (UPat(Ops.AFTER, name="u"), collect_stores),