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.

spec/tinyspec.tex

@@ -109,14 +109,14 @@ A value \op{Call} is void: its \op{Sink} body stores to output \op{Param}s bound
 \end{tabular}
 
 %% ============================================================
-\subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace}
+\subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace, anonymous store}
 
 \begin{tabular}{@{}l l l l@{}}
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Load} & (buf, alt?, gate?) & device, addrspace & Read (pull) from buffer into a new anonymous buffer. \\
-  & & & Note: this replaces \op{Copy} and \op{Contiguous}. \\
+\op{Load} & (buf, alt?, gate?) & --- & Read from buffer into AddrSpace.ALU. \\
+\op{Copy} & (buf,) & device & Copy buf to device arg. \\
 \bottomrule
 \end{tabular}
 

test/amd/test_roundtrip.py

@@ -79,7 +79,7 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
             buf_pool[dst_id] = dst_buf.nbytes
           # Get source data if it's from numpy/CPU
           if hasattr(src_buf, 'base') and src_buf.base is not None and src_buf.base.is_allocated():
-            src_data = bytes(src_buf.base.as_memoryview())
+            src_data = bytes(src_buf.base._buf)
             buf_data[dst_id] = src_data
       elif ast.op is Ops.PROGRAM:
         info = ast.arg

test/backend/test_arange.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
 from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException
+from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException, is_hcq2_device
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
@@ -153,7 +153,7 @@ class TestIndexing(unittest.TestCase):
       GlobalCounters.reset()
       z = emb(x).realize()
       self.assertLessEqual(GlobalCounters.global_ops, op_limit)
-      assert_kernel_count(2)
+      assert_kernel_count(3 if is_hcq2_device() else 2)
     if getenv("CHECK", 1):
       import torch
       with torch.no_grad():

test/backend/test_assign.py

@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad import Device, dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import temp, DEV, Context
-from test.helpers import assert_kernel_count, needs_second_gpu
+from test.helpers import assert_kernel_count, needs_second_gpu, is_hcq2_device
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -43,7 +43,7 @@ class TestAssign(unittest.TestCase):
     # it should copy into the empty buffer
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(1)
+    assert_kernel_count(2 if is_hcq2_device() else 1)
 
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
@@ -619,7 +619,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(3)  # TODO: first copy is dead, could be 2
+    assert_kernel_count(5 if is_hcq2_device() else 3)  # TODO: first copy is dead, could be 2
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_nested_after_contiguous_store_no_init(self):
@@ -629,7 +629,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(1)
+    assert_kernel_count(2 if is_hcq2_device() else 1)
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_assign_temporary_copy_reshape(self):
@@ -637,7 +637,7 @@ class TestAssign(unittest.TestCase):
     c = Tensor.empty(2, 2).assign(a.to(None))
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(1)
+    assert_kernel_count(2 if is_hcq2_device() else 1)
     self.assertEqual(c.tolist(), [[1., 2], [3, 4]])
 
 class TestAssignOrdering(unittest.TestCase):
@@ -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/backend/test_wait_loop.py

@@ -141,7 +141,7 @@ class TestWaitLoop(unittest.TestCase):
 class TestVolatileLoops(unittest.TestCase):
   def test_async_wait_ext(self):
     sig_buf = Buffer(Device.DEFAULT, 1, dtypes.int, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
-    try: sig_view = sig_buf.host.view(fmt='i')
+    try: sig_view = sig_buf.as_memoryview(force_zero_copy=True).cast('i')
     except (AssertionError, NotImplementedError): self.skipTest(f"{Device.DEFAULT} does not support host-visible buffers")
     sig_view[0] = 0
 

test/device/test_hcq2.py

@@ -81,7 +81,7 @@ class TestHCQ2Schedule(unittest.TestCase):
       f(x)
       return f(x), f.captured._linear, [x.uop.base]
     out = chain(x, n)
-    return out, compile_linear(out.schedule_linear(), input_uops=inputs, cache=True), inputs
+    return out, compile_linear(out.schedule_linear(), input_uops=inputs), inputs
 
   def test_jit_has_no_rt_buffers(self):
     dev = Device[Device.DEFAULT]
@@ -116,15 +116,17 @@ class TestHCQ2Schedule(unittest.TestCase):
   def test_host_copies(self):
     dev = Device[Device.DEFAULT]
     if not dev.has_copy_queue: self.skipTest("copy queue required")
-    for host_device in ("CPU", "PYTHON", "NPY", "DISK"):
-      for upload in (False, True):
-        with self.subTest(host_device=host_device, 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),))
-          compiled = compile_linear(linear, profile=False)
-          self.assertEqual(len(compiled.src), 2 if host_device == "DISK" else 1)
-          self.assertEqual(sum(call_is_hcq(call) for call in compiled.src), 1)
+    for host_device in ("CPU", "NPY", "DISK"):
+      for direct in (False, True):
+        for upload in (False, True):
+          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=(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)
+            self.assertEqual(sum(call_is_hcq(call) for call in compiled.src), 1)
 
   def test_large_eager_not_cached(self):
     _, compiled, inputs = self.compiled(65)
@@ -141,7 +143,7 @@ class TestHCQ2Schedule(unittest.TestCase):
           before = tuple(inputs)
           with rt_views() as borrowed:
             for linear in (compiled, linked):
-              self.assertIs(compile_linear(linear, input_uops=inputs, cache=not jit), linear)
+              self.assertIs(compile_linear(linear, input_uops=None if jit else inputs), linear)
           self.assertEqual(tuple(inputs), before)
           self.assertFalse(borrowed)
           run_linear(linked, input_uops=inputs, jit=True, wait=True)
@@ -181,9 +183,20 @@ class TestHCQ2Schedule(unittest.TestCase):
   def test_map_cpu_buffer_preserves_contents(self):
     src = Buffer("CPU", 16, dtypes.uint8, preallocate=True)
     data = bytes(range(16))
-    src.host[:] = data
+    src.as_memoryview(force_zero_copy=True)[:] = data
     src.get_buf(Device.DEFAULT)
-    self.assertEqual(bytes(src.as_memoryview()), data)
+    self.assertEqual(bytes(src.as_memoryview(force_zero_copy=True)), data)
+
+  def test_staged_copy_roundtrip(self):
+    # a host buffer the device cannot read copies in chunks through a small ring of staging slots: every rotation must land bit-exact
+    stage = Buffer("CPU", size:=1 << 16, dtypes.uint8, preallocate=True)
+    for npdt in (np.uint8, np.float32):
+      with self.subTest(dtype=npdt.__name__):
+        n = (size // 2 // np.dtype(npdt).itemsize) * 9 + 7 # nine rotations of a two slot ring, plus a short tail
+        data = np.arange(n, dtype=np.int64).astype(npdt)
+        with patch.object(hcq2, "STAGING_SIZE", size), patch.object(hcq2, "STAGING_SLOTS", 2), patch.object(hcq2, "_staging", lambda: stage):
+          out = Tensor(data).to(Device.DEFAULT).contiguous().realize()
+          np.testing.assert_equal(out.numpy(), data)
 
   def test_rt_patches_are_inputs_and_vars_only(self):
     x = Tensor.rand(17, 33).contiguous().realize()

test/helpers.py

@@ -3,7 +3,6 @@ from dataclasses import replace
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.device import Buffer
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.codegen import to_program
@@ -126,14 +125,13 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
   allocator = dev.allocator
   bufs = []
   for buf_dt, data in inputs or []:
-    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
-    allocator._copyin(buf.buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
+    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize).buf)
+    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp.param(0, uop.dtype, 1)
   prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
   prog = dev.runtime(prg.to_elf())
-  out_buf = Buffer("PYTHON", 1, uop.dtype, preallocate=True)
-  prog(out_buf._buf, *[b.buf for b in bufs], vals=vals)
-  return out_buf.as_memoryview().cast(uop.dtype.fmt or "").tolist()[0]
+  prog(out_buf:=allocator.alloc(uop.dtype.itemsize).buf, *bufs, vals=vals)
+  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 
 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   sink = UOp.group(*u)

test/mockgpu/amd/emu.py

@@ -1886,8 +1886,8 @@ class WaveState:
       ctypes.memset(self.accvgpr_buf._buf, 0, vgpr_size * 4)
     else:
       self.accvgpr_buf = self.vgpr_buf
-    self._vgpr_mv = self.vgpr_buf.host.view(fmt='I').mv
-    self._sgpr_mv = self.sgpr_buf.host.view(fmt='I').mv
+    self._vgpr_mv = self.vgpr_buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('I')
+    self._sgpr_mv = self.sgpr_buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('I')
     # Zero memory using ctypes memset (much faster than Python loops)
     ctypes.memset(self.vgpr_buf._buf, 0, vgpr_size * 4)
     ctypes.memset(self.sgpr_buf._buf, 0, SGPR_COUNT * 4)

test/unit/test_after.py

@@ -42,19 +42,17 @@ class TestAfterCounterexamples(unittest.TestCase):
     # y = x**4, so dy/dx = 4*x**3. Currently raises "cycle detected while indexing".
     self.assertEqual(y.sum().gradient(x)[0].tolist(), [32.])
 
-  @unittest.expectedFailure
   def test_partial_store_gradient(self):
     x = Tensor([2., 3.]).realize()
     y = Tensor(x.uop.after(x[:1].uop.store(4)))
-    # y = [4, x[1]]. Currently returns [0., 0.].
+    # y = [4, x[1]]; only the untouched element depends on x.
     self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
 
-  @unittest.expectedFailure
   def test_partial_store_source_gradient(self):
     x = Tensor([4.])
     y = Tensor([2., 3.]).realize()
     z = Tensor(y.uop.after(y[:1].uop.store(x.uop)))
-    # x contributes once, not twice. Currently returns [2.].
+    # x contributes once, not twice.
     self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
 
   def test_unrelated_store_gradient(self):
@@ -64,14 +62,26 @@ class TestAfterCounterexamples(unittest.TestCase):
     # Zeroing y does not change x.
     self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
 
-  @unittest.expectedFailure
   def test_after_dependency_gradient(self):
     x = Tensor([2., 3.])
     y = x.clone()
     y[:1].assign(0)
-    # View assign creates a nested AFTER; currently raises in backward.
+    # View assign is an AFTER on a partial STORE; only the untouched element depends on x.
     self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
 
+  def test_view_assign_gradient(self):
+    for view, expected in ((lambda t: t.reshape(3, 2)[1:], [[1., 1., 0.], [0., 0., 0.]]),
+                           (lambda t: t.permute(1, 0)[1:], [[1., 0., 0.], [1., 0., 0.]]),
+                           (lambda t: t.flip((0, 1))[:1], [[1., 1., 1.], [0., 0., 0.]])):
+      with self.subTest(expected=expected):
+        x = Tensor([[1., 2., 3.], [4., 5., 6.]])
+        y = x.clone()
+        v = Tensor.full(view(y).shape, 7.)
+        view(y).assign(v)
+        gx, gv = y.sum().gradient(x, v)
+        self.assertEqual(gx.tolist(), expected)
+        self.assertEqual(gv.tolist(), Tensor.ones(v.shape).tolist())
+
   @unittest.expectedFailure
   def test_unordered_overlapping_stores_rejected(self):
     x = Tensor([0.]).realize().uop

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)

test/unit/test_buffer.py

@@ -16,11 +16,10 @@ class TestBuffer(unittest.TestCase):
 
   def test_mapping(self):
     b = Buffer("CPU", 8, dtypes.uint8, initial_value=b"abcdefgh")
-    self.assertEqual(b.get_buf("PYTHON"), b._buf)
+    self.assertIs(b.get_storage("PYTHON").meta, b.get_buf("PYTHON"))
     v = b.view(4, dtypes.uint8, 2)
     mapped = v.get_storage("PYTHON")
-    self.assertEqual(mapped.buf, b._buf + 2)
-    self.assertEqual(bytes(mapped.host.mv), b"cdef")
+    self.assertEqual(bytes(mapped.buf), b"cdef")
     self.assertIs(mapped.host, v.host)
     self.assertIsNone(mapped.meta)
     self.assertIs(v.get_storage("PYTHON"), mapped)
@@ -34,7 +33,7 @@ class TestBuffer(unittest.TestCase):
     self.assertFalse(v.is_allocated())
     v.host[:] = b"test"
     self.assertIsNot(v.get_storage("PYTHON"), old)
-    self.assertEqual(bytes(v.get_storage("PYTHON").host.mv), b"test")
+    self.assertEqual(bytes(v.get_buf("PYTHON")), b"test")
 
   def test_cache_owned_storage_only(self):
     for opaque in (None, memoryview(bytearray(8))):

tinygrad/device.py

@@ -2,11 +2,10 @@ from __future__ import annotations
 from dataclasses import dataclass, replace, field
 from collections import defaultdict
 from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
-import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct, mmap
-from tinygrad.helpers import WIN, mv_address, to_mv, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
+import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct
+from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
-from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up, is_numpy_ndarray
-from tinygrad.helpers import cpu_profile
+from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up
 from tinygrad.dtype import DType, _to_np_dtype
 from tinygrad.runtime.support.memory import MMIOInterface
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
@@ -166,11 +165,6 @@ class Buffer:
       storage = replace(self.base.get_storage(), buf=self.allocator._offset(self.base._buf, self.nbytes, self.offset), maps={})
     elif opaque is not None:
       self.options = replace(self.options, nolru=True)
-      if is_numpy_ndarray(opaque): # readonly arrays stay in place, mv_address can't take their address
-        if not opaque.flags.c_contiguous: opaque = opaque.copy(order='C')
-        opaque = BufferStorage(addr:=opaque.ctypes.data, memoryview(opaque), MMIOInterface(addr, self.nbytes))
-      elif isinstance(opaque, memoryview):
-        opaque = BufferStorage(addr:=mv_address(opaque) if self.nbytes else 0, opaque, MMIOInterface(addr, self.nbytes))
       storage = opaque if isinstance(opaque, BufferStorage) else BufferStorage(opaque)
     else: storage = self.allocator.alloc(self.nbytes, self.options)
     storage = replace(storage, host=storage.host.view(self.offset, self.nbytes, fmt='B') if storage.host is not None else None)
@@ -199,10 +193,7 @@ class Buffer:
     buf:bytearray|pickle.PickleBuffer|None = None
     if self._base is not None:
       return self.__class__, (self.device, self.size, self.dtype, None, None, None, self.base, self.offset, self.is_allocated())
-    if self.device == "NPY": # the array pickles itself, no staging copy
-      import numpy as np
-      arr = np.frombuffer(self.as_memoryview(allow_zero_copy=True), _to_np_dtype(self.dtype))
-      return self.__class__, (self.device, self.size, self.dtype, arr, self.options, None)
+    if self.device == "NPY": return self.__class__, (self.device, self.size, self.dtype, self._buf, self.options, None)
     if self.is_allocated():
       buf = pickle.PickleBuffer(self.as_memoryview()) if protocol >= 5 else bytearray(self.as_memoryview())
     return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf)
@@ -217,10 +208,12 @@ class Buffer:
     if self.is_allocated() and hasattr(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf)
     return None
 
-  def as_memoryview(self, allow_zero_copy=False) -> memoryview:
-    if allow_zero_copy and (mv:=self._host_mv()) is not None:
-      for device in {self.device, *self.base.get_storage().maps}: Device[device].synchronize()
+  def as_memoryview(self, allow_zero_copy=False, force_zero_copy=False, no_sync=False) -> memoryview:
+    # zero copy with as_memoryview (disabled by default due to use after free)
+    if (force_zero_copy or allow_zero_copy) and (mv:=self._host_mv()) is not None:
+      if not no_sync: self.allocator.dev.synchronize()
       return mv
+    assert not force_zero_copy, "force zero copy was passed, but copy is required"
     Buffer("PYTHON", self.size, self.dtype, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self)
     return mv
 
@@ -235,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:
@@ -292,23 +285,6 @@ class Allocator(Generic[DeviceType]):
   # def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
   def _encode_decode(self, bufout, bufin, desc, hist:list, shape:tuple[int,...], frame_pos:int): raise NotImplementedError("need encdec") # optional
 
-class HostAllocator(Allocator):
-  def __init__(self, dev): super().__init__(dev, supports_copy_from_disk=False, supports_transfer=False)
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
-    if options.external_ptr is not None: addr, buf = options.external_ptr, None
-    elif WIN: addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE))
-    else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE))
-    return BufferStorage(addr, buf, MMIOInterface(addr, size, fmt='B'))
-
-  def _copyin(self, dest:int, src:memoryview): # a slice copy takes readonly sources, memmove doesn't
-    self.dev.synchronize()
-    with cpu_profile(f"TINY -> {self.dev.device}", f"{self.dev.device}:COPY"): to_mv(dest, src.nbytes)[:] = src.cast('B')
-  def _copyout(self, dest:memoryview, src:int):
-    self.dev.synchronize()
-    with cpu_profile(f"{self.dev.device} -> TINY", f"{self.dev.device}:COPY"): dest[:] = to_mv(src, dest.nbytes)[:]
-  def _map(self, buf:Buffer) -> BufferStorage: return BufferStorage(buf.host.addr)
-  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
-
 class DepsTracker:
   def __init__(self):
     # tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.

tinygrad/engine/jit.py

@@ -69,7 +69,7 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
   # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
   linear = linear.substitute({u: UOp.param(i, u.dtype, u.max_numel(), u.device) for i,u in enumerate(input_uops)}, walk=True)
   linear = memory_plan_rewrite(linear, held_bufs)
-  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value), input_uops=input_uops, cache=False)
+  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
   if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
   if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View graphed linear")
   return linear

tinygrad/engine/realize.py

@@ -157,9 +157,9 @@ def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
          and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
       dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
     elif src.device.split(":")[0] in HCQ_DEVS and dest._host_mv() is not None:
-      dst_mv, src_mv = dest.as_memoryview(allow_zero_copy=True), src.as_memoryview(allow_zero_copy=True)
+      dst_mv, src_mv = dest.as_memoryview(force_zero_copy=True), src.as_memoryview(force_zero_copy=True)
       with cpu_profile(f"{src.device} -> TINY", f"{src.device}:COPY"): dst_mv[:] = src_mv[:]
-    elif dest._host_mv() is not None: src.allocator._copyout(dest.as_memoryview(allow_zero_copy=True), src._buf)
+    elif dest._host_mv() is not None: src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
     else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
   return []
 
@@ -282,11 +282,11 @@ pm_exec = PatternMatcher([
 
 from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV, HCQInfo, HCQ_DEVS # noqa: E402 # down here, hcq2 imports realize
 
-def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None, cache=False) -> UOp:
+def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
   linear = lower_and_compile(linear)
-  linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile, cache=cache)
+  linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
   return linear
 
 def link_linear(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UOp:
@@ -294,13 +294,13 @@ def link_linear(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) ->
 
 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
   inputs = list(input_uops)
-  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs, cache=True), input_uops=inputs)
+  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs), input_uops=inputs)
   ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2)
   for call in linear.src: track_stats(ctx, call.without_after, perf_counter_us(), pm_exec.rewrite(call.without_after, ctx))
 
 def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> Iterator[float]:
   ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
-  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True, cache=False), allow_cache=ctx.cache)
+  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), allow_cache=ctx.cache)
   while True:
     if clear_l2:
       if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()

tinygrad/helpers.py

@@ -5,8 +5,7 @@ import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfi
 from collections import defaultdict
 import shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
 from dataclasses import dataclass, field, replace
-from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload, TYPE_CHECKING
-if TYPE_CHECKING: import numpy
+from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -36,7 +35,6 @@ def get_shape(x) -> tuple[int, ...]:
   return (len(subs),) + (subs[0] if subs else ())
 def is_image_shape(shape): return shape is not None and len(shape) == 3 and shape[-1] == 4
 def all_int(t: Sequence[Any]) -> TypeGuard[tuple[int, ...]]: return all(isinstance(s, int) for s in t)
-def is_numpy_ndarray(x) -> TypeGuard[numpy.ndarray]: return str(type(x)) == "<class 'numpy.ndarray'>"
 def colored(st, color:str|None, background=False): # replace the termcolor library
   if NO_COLOR: return st
   colors = ['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']

tinygrad/mixin/gradient.py

@@ -2,7 +2,7 @@ from typing import cast
 import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
-from tinygrad.dtype import sum_acc_dtype
+from tinygrad.dtype import dtypes, sum_acc_dtype
 from tinygrad.function import renumber_invalid_outputs
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
@@ -67,6 +67,19 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   ret_set = set(ret_pos)
   return (None,) + tuple(None if i in ret_set else (bwd_outs[gb_map[i]] if i in gb_map else None) for i in range(len(args)))
 
+def partial_store_gradient(ctx:UOp, dest:UOp, view:UOp):
+  # A write through a non-overlapping view replaces only that region of the returned state.
+  path, base = [], view
+  while base is not dest and base.op in {Ops.RESHAPE, Ops.SHRINK, Ops.PERMUTE, Ops.FLIP}:
+    path.append(base)
+    base = base.src[0]
+  if base is not dest: return None
+  grad = ctx
+  for mop in reversed(path): grad = mop.replace(src=(grad,)+mop.src[1:])
+  mask = grad.const_like(1)
+  for mop in path: mask = pm_gradient.rewrite(mop, ctx=mask)[0]
+  return mask.cast(dtypes.bool).where(0, ctx), grad
+
 # ctx is grad_output
 pm_gradient = PatternMatcher([
   (UPat(Ops.CAST, name="ret"), lambda ctx, ret: (ctx.cast(ret.src[0].dtype),)),
@@ -104,6 +117,7 @@ pm_gradient = PatternMatcher([
    lambda ctx, dest, t: (ctx, None) if t.buf_uop is not dest.buf_uop else None),
   # clone/assign gradient passes through to val
   (UPat(Ops.AFTER, src=(UPat(name="dest"), UPat(Ops.STORE, src=(UPat(name="dest"), UPat())))), lambda ctx,dest: (None, ctx)),
+  (UPat(Ops.AFTER, src=(UPat(name="dest"), UPat(Ops.STORE, src=(UPat(name="view"), UPat())))), partial_store_gradient),
   (UPat(Ops.STORE, src=(UPat(), UPat())), lambda ctx: (None, ctx)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),

tinygrad/renderer/llvmir.py

@@ -152,7 +152,7 @@ class LLVMRenderer(Renderer):
 
   extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast
   def _render_fn(self, name:str, args:list[tuple[str,UOp]], kernel:list[str], prefix:list[str]|None=None) -> str:
-    # NOTE: HostAllocator promises 0x20 alignment
+    # NOTE: CPUAllocator promises 0x20 alignment
     sargs = ", ".join([f"{ldt(u.dtype, ptr=u.addrspace == AddrSpace.GLOBAL)}{' noalias align 32' if u.addrspace == AddrSpace.GLOBAL else ''} " + \
       name for name,u in args])
     return "\n".join((prefix or []) + [f"define{' ' + self.abi if self.abi else ''} void @{name}({sargs}) #0", "{"] + kernel + ["  ret void\n}"])

tinygrad/runtime/ops_amd.py

@@ -8,7 +8,7 @@ from tinygrad.uop.ops import sint, UOp, ProgramInfo
 from tinygrad.device import BufferStorage, BufferSpec, Buffer, Device, Allocator, Compiled, ProfileProgramEvent
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, prod, colored
-from tinygrad.helpers import ceildiv, unwrap, pluralize, HCQ2, ContextVar, VIZ
+from tinygrad.helpers import ceildiv, unwrap, pluralize, HCQ2, mv_address, ContextVar, VIZ
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
@@ -651,10 +651,8 @@ class KFDIface:
     if owned: kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=handle)
 
   def map(self, buf:Buffer) -> BufferStorage:
-    if buf.device.split(":")[0] in {"CPU", "PYTHON", "NPY"}:
-      if buf._buf % 0x1000: raise RuntimeError("Host mapping requires a page-aligned address")
+    if buf.device.split(":")[0] == "CPU":
       return replace(mem:=self.alloc(buf.nbytes, host=True, cpu_addr=buf._buf), meta=(mem.meta.handle, True))
-    if buf.device.split(":")[0] != "AMD": raise RuntimeError(f"Cannot map {buf.device} on {self.dev.device}")
     self._map_handle(buf.meta.handle)
     return BufferStorage(buf._buf, (buf.meta.handle, False))
 
@@ -802,7 +800,9 @@ class PCIIface(PCIIfaceBase):
   def device_fini(self): self.dev_impl.fini()
 
 class USBAllocator(AMDAllocator): # the host program reads another device's memory in place: its bytes are the mapping
-  def map(self, buf:Buffer) -> BufferStorage: return BufferStorage(buf.host.addr, buf.host.mv)
+  def map(self, buf:Buffer) -> BufferStorage:
+    mv = buf.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True)
+    return BufferStorage(mv_address(mv), mv)
   def _unmap(self, mapping:BufferStorage): pass
 
 class USBIface(PCIIface):
@@ -891,6 +891,7 @@ class AMDDevice(HCQ2Compiled):
     if self.is_usb: # the submits write the rings over the link, the copies go through the controller's sram (usb.py)
       self.pm_batch, self.pm_lower = pm_usb_batch, pm_usb_lower
       self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
+      self.host_devs = frozenset({"CPU", "NPY", "DISK"}) # the host program streams numpy and files in place
 
     # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
     self.pmc_enabled, self.sqtt_enabled = PROFILE > 0 and PMC > 0, PROFILE > 0 and SQTT > 0

tinygrad/runtime/ops_cl.py

@@ -3,7 +3,7 @@ from typing import cast
 import ctypes, hashlib
 from tinygrad.runtime.autogen import opencl as cl
 from tinygrad.runtime.support import c
-from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, suppress_finalizing, unwrap, round_up, is_image_shape
+from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, mv_address, suppress_finalizing, unwrap, round_up, is_image_shape
 from tinygrad.renderer.cstyle import OpenCLRenderer
 from tinygrad.device import BufferStorage, BufferSpec, Allocator, Compiled, Compiler, CompileError, TinyELF, Program
 
@@ -82,8 +82,9 @@ class CLAllocator(Allocator['CLDevice']):
   @suppress_finalizing
   def _free(self, storage:BufferStorage, options:BufferSpec): check(cl.clReleaseMemObject(storage.buf))
   def _copyin(self, dest:cl.cl_mem, src:memoryview):
-    self.dev.pending_copyin.append(src:=memoryview(bytearray(src))) # NOTE: these can't be freed until the GPU actually executes this command
+    if mv_address(src) % 16: src = memoryview(bytearray(src))
     check(cl.clEnqueueWriteBuffer(self.dev.queue, dest, False, 0, len(src)*src.itemsize, from_mv(src), 0, None, None))
+    self.dev.pending_copyin.append(src)    # NOTE: these can't be freed until the GPU actually executes this command
   def _copyout(self, dest:memoryview, src:cl.cl_mem):
     check(cl.clEnqueueReadBuffer(self.dev.queue, src, False, 0, len(dest)*dest.itemsize, from_mv(dest), 0, None, None))
     self.dev.synchronize()

tinygrad/runtime/ops_cpu.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 import platform, sys, ctypes, mmap, struct, time
 from typing import cast
-from tinygrad.helpers import OSX, WIN, mv_address, suppress_finalizing, unwrap, data64_le
-from tinygrad.device import TinyELF, Program, Device, HostAllocator
+from tinygrad.helpers import to_mv, from_mv, OSX, WIN, mv_address, suppress_finalizing, unwrap, data64_le
+from tinygrad.device import BufferStorage, BufferSpec, TinyELF, Program, Device, Buffer, Allocator
+from tinygrad.runtime.support.memory import MMIOInterface
 from tinygrad.runtime.support.hcq2 import HCQ2Compiled
 from tinygrad.runtime.support.c import DLL
 from tinygrad.renderer.cstyle import ClangRenderer
@@ -73,11 +74,30 @@ class CPUProgram(Program['CPUDevice']):
   def __del__(self):
     if sys.platform == 'win32': ctypes.windll.kernel32.VirtualFree(ctypes.c_void_p(self.addr), ctypes.c_size_t(0), 0x8000) #0x8000 - MEM_RELEASE
 
+class CPUAllocator(Allocator['CPUDevice']):
+  def __init__(self, dev:CPUDevice): super().__init__(dev, supports_copy_from_disk=False, supports_transfer=False)
+  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
+    if options.external_ptr is not None: addr, buf = options.external_ptr, None
+    elif WIN: addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE))
+    else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE))
+    return BufferStorage(addr, buf, MMIOInterface(addr, size, fmt='B'))
+
+  def _copyin(self, dest:int, src:memoryview):
+    self.dev.synchronize()
+    ctypes.memmove(dest, from_mv(src), len(src))
+  def _copyout(self, dest:memoryview, src:int):
+    self.dev.synchronize()
+    dest[:] = to_mv(src, dest.nbytes)[:]
+  def _map(self, buf:Buffer) -> BufferStorage:
+    if not isinstance(host:=buf.get_storage().host, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
+    return BufferStorage(host.addr)
+  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
+
 class CPUDevice(HCQ2Compiled):
   wait_timeout_ms, has_copy_queue = 30000, False
 
   def __init__(self, device:str=""):
-    super().__init__(device, HostAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram,
+    super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram,
       arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
 
   def synchronize(self, timeout:int|None=None): # a host read is safe once every device timeline caught up

tinygrad/runtime/ops_npy.py

@@ -1,4 +1,15 @@
-from tinygrad.device import Compiled, HostAllocator
+import numpy as np
+from tinygrad.helpers import flat_mv
+from tinygrad.device import BufferStorage, MMIOInterface, Compiled, Allocator
+
+class NpyAllocator(Allocator['NpyDevice']):
+  def _alloc(self, size:int, options) -> BufferStorage:
+    return BufferStorage(arr:=np.empty(size, dtype=np.uint8), arr, MMIOInterface(arr.ctypes.data, size))
+
+  def _as_buffer(self, src:np.ndarray) -> memoryview: return flat_mv(np.require(src, requirements='C').data)
+  def _copyout(self, dest:memoryview, src:np.ndarray): dest[:] = self._as_buffer(src)
+  def _offset(self, buf:np.ndarray, size:int, offset:int) -> np.ndarray:
+    return np.require(buf, requirements='C').reshape(-1).view(np.uint8)[offset:offset+size]
 
 class NpyDevice(Compiled):
-  def __init__(self, device:str): super().__init__(device, HostAllocator(self), [], None)
+  def __init__(self, device:str): super().__init__(device, NpyAllocator(self), [], None)

tinygrad/runtime/ops_nv.py

@@ -510,11 +510,9 @@ class NVKIface:
 
   def map(self, buf:Buffer) -> BufferStorage:
     mem = buf.meta
-    if buf.device.split(":")[0] in {"CPU", "PYTHON", "NPY"}:
-      if buf._buf % 0x1000: raise RuntimeError("Host mapping requires a page-aligned address")
+    if buf.device.split(":")[0] == "CPU":
       if (mem:=next((m.meta[0] for d, m in buf.get_storage().maps.items() if d.startswith("NV")), None)) is None:
         return replace(mem:=self.alloc(buf.nbytes, host=True, cpu_addr=buf._buf), meta=(mem.meta, True))
-    elif buf.device.split(":")[0] != "NV": raise RuntimeError(f"Cannot map {buf.device} on {self.dev.device}")
     return replace(mapping:=self._gpu_uvm_map(buf._buf, mem.length, mem.hMemory, create_range=False), meta=(mapping.meta, False))
 
   def _alloc_gpu_vaddr(self, size, alignment=(4 << 10), force_low=False):

tinygrad/runtime/ops_python.py

@@ -6,8 +6,8 @@ from typing import Any, TYPE_CHECKING
 import pickle, base64, itertools, time, sys, functools, ctypes
 from dataclasses import replace
 from tinygrad.dtype import bitcast, DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
-from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, to_mv, mv_address
-from tinygrad.device import HostAllocator, Compiled, Compiler, Program, TinyELF
+from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile, mv_address
+from tinygrad.device import BufferStorage, MMIOInterface, Buffer, Compiled, Compiler, Allocator, Program, TinyELF
 from tinygrad.renderer import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer
@@ -55,7 +55,7 @@ class PythonProgram(Program['PythonDevice']):
     warp_size = len(warp)
     for idxs in itertools.product(*[range(x) for x in global_size[::-1]]):
       values: dict[UOp, Any] = {}
-      pbufs: list[int] = list(bufs)
+      pbufs: list[memoryview] = list(bufs)
       pvals: list[int] = list(vals)
       exec_masks = [[True] * warp_size]
       i = 0
@@ -101,8 +101,7 @@ class PythonProgram(Program['PythonDevice']):
             # REGs are per thread
             values[u] = [memoryview(bytearray(u.max_numel()*u.dtype.itemsize)).cast(storage_fmt) for _ in range(warp_size)]
           else:
-            size = u.max_numel() * u.dtype.itemsize
-            buf = memoryview(bytearray(size)) if u.op is not Ops.PARAM else to_mv(pbufs.pop(0), size)
+            buf = memoryview(bytearray(u.max_numel()*u.dtype.itemsize)) if u.op is not Ops.PARAM else pbufs.pop(0)
             values[u] = [buf.cast(storage_fmt)] * warp_size
         elif u.op is Ops.SPECIAL:
           if u.arg[0] == 'g': values[u] = [idxs[2-int(u.arg[-1])]] * warp_size
@@ -237,6 +236,18 @@ class PythonRenderer(Renderer):
 
   def supported_dtypes(self): return {d for d in super().supported_dtypes() if d != dtypes.half or sys.version_info >= (3, 12)}
 
+class PythonAllocator(Allocator['PythonDevice']):
+  def _alloc(self, size:int, options) -> BufferStorage:
+    return BufferStorage(buf:=memoryview(bytearray(size)), buf, MMIOInterface(mv_address(buf), size))
+
+  def _as_buffer(self, src) -> memoryview: return src
+  def _copyin(self, dest, src:memoryview):
+    with cpu_profile("TINY -> PYTHON", f"{self.dev.device}:COPY"): dest[:] = src
+  def _copyout(self, dest:memoryview, src):
+    with cpu_profile("PYTHON -> TINY", f"{self.dev.device}:COPY"): dest[:] = src
+  def map(self, buf:Buffer) -> BufferStorage: return BufferStorage(mv:=buf.as_memoryview(force_zero_copy=True), mv)
+  def _offset(self, buf:memoryview, size:int, offset:int): return buf[offset:offset+size]
+
 class PythonDevice(Compiled):
   def __init__(self, device:str):
-    super().__init__(device, HostAllocator(self), [PythonRenderer], PythonProgram)
+    super().__init__(device, PythonAllocator(self), [PythonRenderer], PythonProgram)

tinygrad/runtime/support/hcq2.py

@@ -11,7 +11,7 @@ from tinygrad.dtype import dtypes, DType, DTYPES_DICT, AddrSpace
 from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.engine.realize import get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear
-from tinygrad.engine.realize import lower_and_compile, _resolve
+from tinygrad.engine.realize import lower_and_compile
 
 # *****************
 # 0. helpers
@@ -41,6 +41,7 @@ def get_enqueue_devs(call:UOp) -> Any|None:
   devs = min(bufs, key=lambda b: not all_devices_in(b.device, HCQ_DEVS)).device
   if not all_devices_in(devs, HCQ_DEVS): return None
   dev = cast(HCQ2Compiled, Device[to_tuple(devs)[0]])
+  if not all(all_devices_in(b.device, HCQ_DEVS | dev.host_devs) for b in bufs): return None
   # a device without a copy queue leaves copies to its allocator
   return devs if call.src[0].op is not Ops.COPY or dev.has_copy_queue else None
 
@@ -102,7 +103,35 @@ def replace_buffer(ctx:tuple[bool, list[UOp], dict[UOp, int]], b:UOp) -> UOp:
 pm_replace_buffers = PatternMatcher([(UPat(Ops.BUFFER, name="b"), replace_buffer)])
 
 # *****************
-# 1.1. prep: unwrap multi
+# 1.1. prep: staging copies
+
+STAGING_SIZE, STAGING_SLOTS = (4 if DEV.interface.startswith("MOCK") else 128) << 20, 2
+
+@functools.cache
+def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preallocate=True)
+
+def _need_staging(a, b):
+  if not all_devices_in(a.device, HCQ_DEVS): return False
+  dev = cast(HCQ2Compiled, Device[to_tuple(a.device)[0]])
+  return not all_devices_in(b.device, HCQ_DEVS | dev.host_devs) and dev.has_copy_queue
+
+def stage_copy(dst:UOp, src:UOp) -> UOp|None:
+  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
+
+  base, it, copies = UOp.from_buffer(_staging()), src.dtype.itemsize, []
+  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 += [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([
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
+])
+
+# *****************
+# 1.2. prep: one call per device: the args pick their lane, the DEVICE axis binds to it
 
 def unwrap_call(call:UOp) -> UOp|None:
   if get_enqueue_devs(call) is None or (n:=max(len(to_tuple(a.device)) for a in get_call_arg_uops(call))) == 1: return None
@@ -111,32 +140,6 @@ def unwrap_call(call:UOp) -> UOp|None:
                                    for i in range(n)))
 pm_unwrap_multi = PatternMatcher([(UPat(Ops.CALL, name="call"), unwrap_call)])
 
-# *****************
-# 1.2. prep: staging copies
-
-STAGING_SIZE, STAGING_SLOTS = (4 if DEV.interface.startswith("MOCK") else 128) << 20, 2
-
-@functools.cache
-def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preallocate=True)
-
-def stage_copy(ctx:tuple[UOp, ...], call:UOp, dst:UOp, src:UOp) -> UOp|None:
-  if (device:=get_enqueue_devs(call)) is None: return None
-  try:
-    for b in (dst, src): cast(Buffer, _resolve(b, ctx).buffer).get_buf(device)
-    return None
-  except (RuntimeError, OSError): _staging().get_buf(device)
-
-  base, it, copies = UOp.from_buffer(_staging()), src.dtype.itemsize, []
-  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)]
-  return UOp(Ops.LINEAR, src=tuple(copies))
-
-pm_insert_copy_staging = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call", allow_any_len=True), stage_copy),
-])
-
 # *****************
 # 2. deps
 
@@ -417,20 +420,20 @@ pm_encode = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.SINK),), name="call",
 
 hcq_compile_cache:dict[tuple[UOp, bool], UOp] = {} # eager templates: a buffer-free linear (uops are hash-consed) to its compiled form
 
-@rewrite_group(lambda linear,input_uops,profile,ret,cache=False: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
-def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool, cache=False) -> UOp:
+@rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
+def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
   if any(isinstance(getattr(c.without_after.arg, "aux", None), HCQInfo) for c in linear.src): return linear # compiled already
-  if input_uops is None: input_uops = []
 
-  if cache:
+  if input_uops is not None:
     use_rt = len(linear.src) < HCQ_CACHE_THRESH # small schedules use runtime address patches so linked schedules can be cached without input buffers
     slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))}
     linear = graph_rewrite(linear, pm_replace_buffers, ctx=(use_rt, input_uops, slots), walk=True, name="replace buffers")
-  linear = graph_rewrite(linear, pm_unwrap_multi+pm_insert_copy_staging+pm_flatten_linear, ctx=tuple(input_uops), name="prep calls")
-  if cache and (cached:=hcq_compile_cache.get(key:=(linear, profile))) is not None: return cached
-  lin = graph_rewrite(sched_batches(linear, profile), pm_encode, walk=True, name="encode")
+    if (cached:=hcq_compile_cache.get(key:=(linear, profile))) is not None: return cached
+  lin = graph_rewrite(linear, pm_unwrap_multi+pm_insert_copy_staging+pm_flatten_linear, name="prep calls")
+  lin = sched_batches(lin, profile)
+  lin = graph_rewrite(lin, pm_encode, walk=True, name="encode")
   with Context(EMULATED_DTYPES=""): final_linear = lower_and_compile(lin)
-  if cache and final_linear is not linear: hcq_compile_cache[key] = final_linear
+  if input_uops is not None and final_linear is not linear: hcq_compile_cache[key] = final_linear
   return final_linear
 
 # *****************
@@ -511,6 +514,7 @@ class HCQ2Compiled(Compiled):
   wait_timeout_ms: float = 30000.0
   sleep_timeout_ms: int|None = None
   rt_nbytes: int = 64 << 20 # the pool every per-linear buffer is carved out of
+  host_devs: frozenset[str] = frozenset({"CPU"})
   pm_encode: PatternMatcher = PatternMatcher([]) # the backend's own encode rules, matched by its submit names
   var_vals: dict[str, int] = {}
 

tinygrad/runtime/support/system.py

@@ -289,11 +289,8 @@ class PCIIfaceBase:
     return [(p + self.pci_dev.bar_info(self.vram_bar)[0], sz) for p, sz in paddrs], AddrSpace.SYS
 
   def map(self, b:Buffer) -> BufferStorage:
-    if b.device.split(":")[0] in {"CPU", "PYTHON", "NPY"}:
-      if b._buf % 0x1000: raise RuntimeError("Host mapping requires a page-aligned address")
-      lo, size = b._buf, round_up(b.nbytes, 0x1000)
-      if not self.dev_impl.mm.va_base <= lo < lo + size <= self.dev_impl.mm.va_base + (1 << self.dev_impl.mm.va_bits):
-        raise RuntimeError(f"Host address {lo:#x} is outside the GPU virtual address range")
+    if b.device.split(":")[0] == "CPU":
+      lo, size = b._buf & ~0xfff, round_up(b._buf + b.nbytes, 0x1000) - (b._buf & ~0xfff)
       System.lock_memory(lo, size)
       paddrs, aspace, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(lo, size)], AddrSpace.SYS, True, True
     elif isinstance(ifa:=getattr(Device[b.device], "iface", None), PCIIfaceBase):

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

@@ -2,10 +2,10 @@
 from __future__ import annotations
 import time, functools, sys, inspect, pathlib, hashlib, weakref
 from dataclasses import dataclass, field, replace
-from typing import Any, Callable, cast, get_args, ParamSpec, TypeVar, Generic, TYPE_CHECKING
+from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, _from_np_dtype, _to_np_dtype, PyConst, AddrSpace
-from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey, is_numpy_ndarray
+from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc, VIZ, pluralize, SPEC
 from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike, UPat, PatternMatcher, GroupOp, graph_rewrite, rewrite_group
 from tinygrad.uop.ops import resolve_returned_after, remove_all_tags
@@ -43,8 +43,10 @@ def creation_copy_is_realized(u:UOp):
 # CONTIGUOUS and AFTER + parents are the only nodes that get updated
 add_tags = PatternMatcher([
   (UPat(Ops.COPY, name="u"), creation_copy_is_realized),
-  # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
+  # no tag on copies that fill an AFTER's whole dest via STORE: merge COPY tag into AFTER (the copy reads that storage).
+  # a partial STORE keeps the tag: the copy mints its own storage like any bare creation copy
+  (UPat(Ops.AFTER, src=(UPat(name="dest"),
+    UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
    lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
   (UPat(Ops.AFTER, name="x"), tag_uop),
   (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(x) if x in ctx.bases else None),
@@ -219,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)
@@ -254,6 +255,8 @@ def _tensor_holds(u:UOp) -> bool: return any((t:=tref()) is not None and t.uop i
 
 # **** Tensor helper functions ****
 
+def is_numpy_ndarray(x) -> "TypeGuard[numpy.ndarray]": return str(type(x)) == "<class 'numpy.ndarray'>"
+
 def _fromnp(x: 'numpy.ndarray') -> UOp:
   ret = UOp.new_buffer("NPY", x.size, _from_np_dtype(x.dtype))
   # fake realize
@@ -430,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}")
 
@@ -446,15 +450,18 @@ class Tensor(RandMixin):
       self.uop = (x.uop.src[0] if x.uop.op is Ops.CONTIGUOUS else x.uop).clone()
       return self
     # STORE+AFTER: STORE is the write effect (void), AFTER wraps the view for correct shape/ranging
-    assign = self.uop.after(self.uop.store(x.uop))
+    assign = self.uop.after(store := self.uop.store(x.uop))
     ib = self.uop
     while ib.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH} and not (ib.has_buffer_identity() and _tensor_holds(ib)): ib = ib.src[0]
     if ib is not self.uop:
       # a partial write needs storage to land in: a pending value gets explicit storage (a clone)
       target = ib if ib.has_buffer_identity(after_ok=True) else ib.clone()
-      if target is not ib: assign = assign.substitute({ib: target}, walk=True)
-      # view assign: replace the node under the views (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
-      _apply_map_to_tensors({ib: target.after(assign)}, name="Embed View Assign")
+      if target is not ib:
+        assign = assign.substitute({ib: target}, walk=True)
+        store = assign.src[1]
+      # view assign: the base reads "after the store into the view" (one AFTER level). replace the node under the
+      # views (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
+      _apply_map_to_tensors({ib: target.after(store)}, name="Embed View Assign")
     else:
       # simple assign
       self.uop = assign
@@ -465,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()
@@ -543,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)
@@ -831,14 +833,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       assert bdtype.fmt is not None, f"{bdtype=} has None fmt"
       ret = UOp.empty(shape:=get_shape(x), dtype=bdtype, device="PYTHON")
       data = struct.pack(f"{prod(shape)}{bdtype.fmt}", *[truncate[bdtype](bdtype.const(xi)) for xi in fully_flatten(x)])
-    if not data: ret.buffer.allocate(memoryview(bytearray()))
-    else: ret.buffer.ensure_allocated().host[:] = data
+    ret.buffer.allocate(memoryview(bytearray(data))) # fake realize. buffer storage must be writable, and bytes isn't
     if ret.dtype != dtype: ret = ret.cast(dtype)
     return ret if ret.device == device else ret.copy_to_device(device)
   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:

tinygrad/viz/serve.py

@@ -49,7 +49,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
                Ops.INDEX: "#CEF9B7", Ops.STACK: "#D8F9E4",
                Ops.WMMA: "#efefc0", Ops.UNSHARD: "#f6ccff", Ops.INS: "#eec4ff",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
-               Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
+               Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#ff90c0", Ops.CUSTOM_FUNCTION: "#bf71b6",
                Ops.CALL: "#00B7C8", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040",
                Ops.LINEAR: "#7DF4FF",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",