diff --git a/test/runtime/test_setitem_schedule.py b/test/runtime/test_setitem_schedule.py index e812c046420de..598705eb65c21 100644 --- a/test/runtime/test_setitem_schedule.py +++ b/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() diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 041bf7a5254b9..c5d1c318829e1 100644 --- a/tinygrad/schedule/__init__.py +++ b/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),