cleanup COPY: it's an anonymous store

* tensor.py and mixins always create COPY for device copies, except DISK
  where a copy is always a store (clone: AFTER+STORE of the disk buffer).
  copy_to_device handles the DISK case so all entry points create COPY.
* callify keeps rewriting copies from creation devices to explicit STOREs
  so they persist; every other COPY stays a COPY until prepare.
* prepare.py: drop the vestigial pm_copy_to_store machinery (buffer
  identity hack, flatten, existing-buf path) and the reshape-on-copy rule.
  A COPY in src[1] of a plain STORE is just removed (a STORE to a buffer
  on a different device is a COPY), and a bare COPY is realized as an
  anonymous store into a fresh call-local buffer on the copy device,
  exactly like contiguous.
* the SDMA copy detector now matches N-dimensional pure copies (identical
  index ranges on both sides) instead of only flattened ones, and the
  source of a cross device STORE is materialized on its own device first.
* runtime paths that really need the op (Buffer.copy_from, hcq staging)
  construct Ops.COPY directly.
* _buffer (numpy path) uses clone instead of contiguous.
* cross device assign works now: the device mismatch error is gone, the
  value is copied to the target device and stored there.

test/backend/test_assign.py

@@ -1255,5 +1255,48 @@ class TestMultiAssign(unittest.TestCase):
       f(out, vi.bind(i))
     self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
 
+class TestCrossDeviceAssign(unittest.TestCase):
+  # CPU:0 and CPU:1 are always available, (Device.DEFAULT, CPU) is a real cross-device pair on GPU runners
+  pairs = [("CPU:0", "CPU:1"), (Device.DEFAULT, "CPU")]
+
+  def test_cross_device_assign(self):
+    for dst_dev, src_dev in self.pairs:
+      a = Tensor.zeros(4, 4, device=dst_dev).realize()
+      a.assign(Tensor.full((4, 4), 3.0, device=src_dev).realize())
+      np.testing.assert_allclose(a.numpy(), np.full((4, 4), 3.0))
+      # the buffer did not move
+      self.assertEqual(a.uop.device, Tensor.empty(1, device=dst_dev).uop.device)
+
+  def test_cross_device_assign_is_copy(self):
+    # a cross device assign is a single COPY call
+    for dst_dev, src_dev in self.pairs:
+      a = Tensor.zeros(5, device=dst_dev).realize()
+      src = Tensor([0.,1.,2.,3.,4.], device=src_dev).realize()
+      linear = Tensor.schedule_linear(a.assign(src))
+      copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
+      sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
+      self.assertEqual(len(copies), 1)
+      self.assertEqual(len(sinks), 0)
+      copy_dst, copy_src = copies[0].src[1:]
+      self.assertEqual(copy_dst.device, a.uop.device)
+      self.assertEqual(copy_src.device, src.uop.device)
+      from tinygrad.engine.realize import run_linear
+      run_linear(linear)
+      np.testing.assert_allclose(a.numpy(), np.arange(5))
+
+  def test_cross_device_assign_unrealized(self):
+    # the source is materialized on its own device before the copy
+    for dst_dev, src_dev in self.pairs:
+      a = Tensor.zeros(8, device=dst_dev).realize()
+      a.assign(Tensor.ones(8, device=src_dev) * 2)
+      np.testing.assert_allclose(a.numpy(), np.full((8,), 2.0))
+
+  def test_cross_device_assign_view(self):
+    # a partial (view) store across devices copies to the target device first
+    for dst_dev, src_dev in self.pairs:
+      a = Tensor.zeros(8, device=dst_dev).realize()
+      a[2:6].assign(Tensor([0.,1.,2.,3.], device=src_dev).realize())
+      np.testing.assert_allclose(a.numpy(), np.array([0, 0, 0, 1, 2, 3, 0, 0], dtype=np.float32))
+
 if __name__ == "__main__":
   unittest.main()

test/device/test_hcq2.py

@@ -122,7 +122,7 @@ class TestHCQ2Schedule(unittest.TestCase):
           with self.subTest(host_device=host_device, direct=direct, upload=upload):
             host, gpu = UOp.new_buffer(host_device, 4, dtypes.uint8), UOp.new_buffer(dev.device, 4, dtypes.uint8)
             src, dst = (host, gpu) if upload else (gpu, host)
-            linear = UOp(Ops.LINEAR, src=(src.copy_to_device(dst.device).call(dst, src),))
+            linear = UOp(Ops.LINEAR, src=(UOp(Ops.COPY, src=(src,), arg=dst.device).call(dst, src),))
             with patch.object(dev, "host_devs", frozenset({"CPU", host_device}) if direct else frozenset({"CPU"})):
               compiled = compile_linear(linear, profile=False)
             self.assertEqual(len(compiled.src), 1 if direct or host_device == "CPU" else 2)

test/unit/test_allreduce.py

@@ -33,8 +33,8 @@ class TestRingAllReduce(unittest.TestCase):
       # N*(N-1) copies for input and output
       copy_count = N*(N-1)*2
       if len(copies) != copy_count: raise KernelCountException(copy_count, len(copies))
-      # N*(N-1) shrinks from other devices becoming contigs, N ALU, N extra contig, reassembly (cat), and mul
-      sink_count = (N*(N-1))+(N)+(N)+(1)+(1)
+      # local stores on the receiving lanes (copies read shard views directly now), partial sums, reassembly (cat), and mul
+      sink_count = (N*(N-1))-(N-1)+1+(N-1)+(1)+(1)
       if len(sinks) != sink_count: raise KernelCountException(sink_count, len(sinks))
       # correctness
       run_linear(linear, var_vals)

tinygrad/device.py

@@ -228,7 +228,7 @@ class Buffer:
     from tinygrad.engine.realize import run_linear
     from tinygrad.uop.ops import UOp, Ops
     du, su = UOp.from_buffer(self), UOp.from_buffer(src)
-    run_linear(UOp(Ops.LINEAR, src=(su.param_like(1).copy_to_device(self.device).call(du, su),)), update_stats=False)
+    run_linear(UOp(Ops.LINEAR, src=(UOp(Ops.COPY, src=(su.param_like(1),), arg=self.device).call(du, su),)), update_stats=False)
     return self
 
   def view(self, size:int, dtype:DType, offset:int) -> Buffer:

tinygrad/runtime/support/hcq2.py

@@ -122,7 +122,8 @@ def stage_copy(dst:UOp, src:UOp) -> UOp|None:
   chunk = (STAGING_SIZE // STAGING_SLOTS) // it
   for i, off in enumerate(range(0, src.max_numel(), chunk)):
     stage = base[(so:=(i % STAGING_SLOTS) * chunk * it):so + (n:=min(chunk, src.max_numel() - off)) * it]
-    copies += [src[off:off+n].copy_to_device("CPU").call(stage, src[off:off+n]), stage.copy_to_device(dst.device).call(dst[off:off+n], stage)]
+    copies += [UOp(Ops.COPY, src=(src[off:off+n],), arg="CPU").call(stage, src[off:off+n]),
+               UOp(Ops.COPY, src=(stage,), arg=dst.device).call(dst[off:off+n], stage)]
   return UOp(Ops.LINEAR, src=tuple(copies))
 
 pm_insert_copy_staging = PatternMatcher([

tinygrad/schedule/__init__.py

@@ -152,8 +152,10 @@ def assert_all_same_devices(ast:UOp):
   devices = dedup([x.device for x in ast.toposort() if x.op is Ops.PARAM and x.device is not None])
   if len(devices) >= 2: raise RuntimeError(f"all buffers must be on the same device: {devices}")
 
-def copy_kernel_to_copy_uop(call:UOp, dst:UOp, src:UOp, r:UOp|None=None):
+def copy_kernel_to_copy_uop(call:UOp, dst:UOp, src:UOp, di:UOp|None=None, si:UOp|None=None, ends:UOp|None=None):
   if dst.device == src.device and not (isinstance(dst.device, str) and dst.device.startswith("DISK")): return None
+  # both sides must be indexed by exactly the same ranges/positions (a pure elementwise copy)
+  if di is not None and si is not None and (di.src[1:] != si.src[1:] or (ends is not None and ends.src[1:] != di.src[1:])): return None
   return call.replace(src=(UOp(Ops.COPY, src=(src,), arg=dst.device),) + call.src[1:])
 
 def simplify_copy_kernel(call:UOp, ast:UOp, dst:UOp, src:UOp):
@@ -170,11 +172,11 @@ pm_copy_from_store = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"), UPat.var("dst"), UPat.var("src")), name="call"), simplify_copy_kernel),
 
   # replace this with a copy if it's a copy
-  (UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(UPat(Ops.CONST, arg=0))
-                .store(UPat(Ops.PARAM, name="src").index(UPat(Ops.CONST, arg=0))).sink(),),
+  (UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(name="di", allow_any_len=True)
+                .store(UPat(Ops.PARAM, name="src").index(name="si", allow_any_len=True)).sink(),),
                 name="call", allow_any_len=True), copy_kernel_to_copy_uop),
-  (UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(UPat(Ops.RANGE, name="r"))
-                .store(UPat(Ops.PARAM, name="src").index(UPat(Ops.RANGE, name="r"))).end(UPat(Ops.RANGE, name="r")).sink(),),
+  (UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(name="di", allow_any_len=True)
+                .store(UPat(Ops.PARAM, name="src").index(name="si", allow_any_len=True)).end(name="ends", allow_any_len=True).sink(),),
                 name="call", allow_any_len=True), copy_kernel_to_copy_uop),
 
   # if it wasn't copy, it currently can't be cross device

tinygrad/schedule/indexing.py

@@ -34,6 +34,9 @@ def realize_srcs(ctx:IndexingContext, rb:UOp) -> None:
 def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp):
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
   if dest.base in src.toposort(enter_calls=False): ctx.realize_map[src] = None
+  # the source of a cross device STORE is materialized on its own device first: the STORE itself is the copy
+  if src.device is not None and dest.device != src.device and not src.has_buffer_identity(after_ok=True):
+    ctx.realize_map[src] = ctx.non_removable[src] = None
 
 def realize_custom_kernel_srcs(ctx:IndexingContext, c:UOp) -> None:
   for s in c.src[1:]:
@@ -49,7 +52,7 @@ pm_generate_realize_map = PatternMatcher([
   (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
   # realize srcs of these
   (UPat((Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
-  # sometimes we need to realize the src of STORE if there's a self-access
+  # sometimes we need to realize the src of STORE if there's a self-access, or if it's a cross device store
   (UPat(Ops.STORE, src=(UPat.var("dest"), UPat.var("src"))), realize_store_after_src),
 ])
 

tinygrad/schedule/prepare.py

@@ -1,4 +1,3 @@
-import itertools
 from tinygrad.dtype import dtypes, to_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp
 from tinygrad.uop.ops import graph_rewrite, rewrite_group, identity_element, resolve_returned_after
@@ -131,6 +130,11 @@ def expand_bitcast(bc:UOp) -> UOp|None:
   parts = [tmp>>8*i*ns for i in range(os//ns)]
   return parts[0].stack(*parts[1:], dim=-1).flatten(-2).cast(new_uint).bitcast(bc.dtype)
 
+def copy_to_anon_store(x:UOp, copy:UOp):
+  # the buffer created here is inside the call and is not persisted, like the buffers created for contiguous
+  buf = UOp.new_buffer(copy.device, prod(x.max_shape), copy.dtype).reshape(x.max_shape)
+  return buf.after(buf.store(x)).reshape(copy.shape)
+
 earliest_rewrites = mop_cleanup+PatternMatcher([
   # resolve calls with RETURNED inputs (inline the body)
   (UPat(Ops.CALL, name="c"), lambda c: resolve_function(c) if c.has_unbound_outputs else None),
@@ -155,8 +159,12 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # copy to same device is a no-op
   (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), lambda x,copy: x if x.device == copy.device else None),
 
-  # copy on reshape is reshape on copy
-  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)),
+  # a COPY in src[1] of a plain STORE can just be removed: a STORE to a buffer on a different device is a COPY
+  (UPat(Ops.STORE, src=(UPat.var("dst"), UPat(Ops.COPY, src=(UPat.var("x"),), name="cpy"))),
+   lambda dst,x,cpy: dst.store(x) if dst.device == cpy.device and dst.has_buffer_identity(after_ok=True) else None),
+
+  # a bare COPY is an anonymous store: realize it as a STORE into a fresh call-local buffer on the copy device
+  (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), copy_to_anon_store),
 
   # reshaping on STORE can be a NOOP
   (UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
@@ -193,31 +201,10 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.AFTER, name="s"), lambda s: s.replace(src=(s.src[0],)+tuple(walk_mop(u) for u in s.src[1:] if u.op is not Ops.NOOP))),
 ])
 
-def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
-  input_src = copy.src[0]
-  # if it's a COPY, we need to give the input buffer identity
-  if not input_src.has_buffer_identity(after_ok=True) and copy.op is Ops.COPY: input_src = input_src.contiguous()
-  input_src = input_src.flatten()
-  if existing_buf is not None:
-    # if the existing buffer is not a full buffer, we can't use it
-    if not existing_buf.has_buffer_identity(after_ok=True): return None
-    # if there's already a buffer, we just use it
-    return existing_buf.flatten().store(input_src)
-  # create the output buffer
-  buf = UOp.new_buffer(copy.device, prod(input_src.max_shape), copy.dtype)
-  # reshape back to input
-  return buf.reshape(input_src.max_shape).after(buf.store(input_src)).reshape(copy.shape)
-
-pm_copy_to_store = PatternMatcher([
-  (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
-  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
-])
-
 @rewrite_group(new_ctx=False)
 def prepare_rangeify(sink:UOp) -> UOp:
   # prepare for rangeify
   tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
   if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
-  tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
   return tsink

tinygrad/tensor.py

@@ -221,11 +221,10 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx=ctx, name="early transform tensor graph")
 
   # collect the stores (never entering call bodies) and map tagged AFTERs to their storage; tags are stripped at the end
-  # copies to disk are stores to the disk buffer; bound Variables are call inputs and RETURNEDs are call outputs
+  # copies to disk are explicit stores to the disk buffer; bound Variables are call inputs and RETURNEDs are call outputs
   # AFTERs on unbound STORAGE (clones) are collected too: the clone's own buffer is the storage, no fresh copy
   for u in big_sink.toposort(enter_calls=False):
-    if (u.op is Ops.COPY and on_disk(u)) or (u.op is Ops.AFTER and not u.is_bound_var and
-        (not u.src[0].unsharded_base.is_unbound or u.src[1].op is Ops.STORE)):
+    if u.op is Ops.AFTER and not u.is_bound_var and (not u.src[0].unsharded_base.is_unbound or u.src[1].op is Ops.STORE):
       ctx.stores.append(u)
       if u.tag: ctx.buffer_map.update({t:graph_rewrite(u.src[0], pm_drop_after).shrink_to(t.shape) for t in u.tag})
   ret = graph_rewrite(UOp.sink(*ctx.stores), pm_replace_buf+remove_all_tags, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements)
@@ -434,8 +433,9 @@ class Tensor(RandMixin):
     x = x._broadcast_to(self.shape)
     if x.dtype in dtypes.weaks: x = x.cast(least_upper_dtype(self.dtype, x.dtype))
     if x.dtype != self.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
+    # an assign is just a STORE: a STORE to a buffer on a different device is a COPY, send the value over first
     if not is_disk and x.uop.device is not None and self.device is not None and self.device != x.device:
-      raise RuntimeError(f"assign device mismatch {self.device} != {x.device}")
+      x = Tensor(x.uop.copy_to_device(self.device))
     if isinstance(self.device, tuple) and x.uop.device is not None and self.uop.axis != x.uop.axis:
       raise RuntimeError(f"multi axis mismatch {self.uop.axis} != {x.uop.axis}")
 
@@ -472,8 +472,12 @@ class Tensor(RandMixin):
     if capturing and not getenv("UNSAFE_ALLOW_JIT_BUFFER"):
       from tinygrad.engine.jit import JitError
       raise JitError("cannot access tensor data during JIT capture, the value will be baked in")
-    x = self.contiguous()
-    if self.uop.device is None or isinstance(self.device, tuple): x = x.clone("CPU")
+    # a named global buffer is needed to read the data out: clone creates one if this value doesn't already have one.
+    # multi device values are materialized per device before gathering to CPU, disk tensors read lazily on allocation
+    if isinstance(self.device, tuple): x = self.clone().clone("CPU")
+    elif self.uop.device is None: x = self.clone("CPU")
+    elif not on_disk(self.uop) and not self.uop.has_buffer_identity(after_ok=True): x = self.clone()
+    else: x = self
     return cast(Buffer, x.realize().uop.buffer).ensure_allocated()
 
   def _data(self) -> memoryview: return self._buffer().as_memoryview()
@@ -550,9 +554,8 @@ class Tensor(RandMixin):
     """
     if self.uop.device is None: return self
     if (device:=canonicalize_device(device)) == self.device: return self
-    # a copy to disk wants to persist, so it inserts a clone: the disk buffer is the storage of the copied value
-    if isinstance(device, str) and device.startswith("DISK"): ret = Tensor(self.uop.clone(device))
-    else: ret = Tensor(self.uop.copy_to_device(device))
+    # a copy to disk is always a store (copy_to_device handles this), all other copies stay COPY until the scheduler
+    ret = Tensor(self.uop.copy_to_device(device))
     if self.grad is not None: ret.grad = self.grad.to(device)
     return ret.is_param_(self.is_param)
 

tinygrad/uop/ops.py

@@ -732,6 +732,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   def copy_to_device(self, device:str|tuple[str, ...], arg=None):
     assert arg is None or isinstance(self.device, tuple)
+    # a copy to a DISK device is always a store: the disk buffer is the storage of the copied value
+    if isinstance(device, str) and device.startswith("DISK"): return self.clone(device)
     inp = self if arg is None else UOp(Ops.MSELECT, src=(self,), arg=arg)
     if inp.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {inp.dtype}")
     return UOp(Ops.COPY, src=(inp.pad_to(inp.max_shape),), arg=device).shrink_to(inp.shape)
@@ -837,7 +839,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def clone(self, device=None) -> UOp:
     device = device or self.device
     ret = self.empty_like(device=device)
-    src = self if self.device is None or self.device == device else self.copy_to_device(device)
+    # a clone to DISK is the store itself (no COPY inside the STORE), a cross device clone stores a COPY
+    src = self if self.device is None or self.device == device or (isinstance(device, str) and device.startswith("DISK")) \
+         else self.copy_to_device(device)
     return ret.after(ret.store(src.cast(ret.dtype)))
   @recursive_property
   def device(self) -> str|tuple[str, ...]|None: