fix view assign replacing prior readers

test/backend/test_assign.py

@@ -327,6 +327,211 @@ def test_assign_corealize_order_independent(self):
       self.assertEqual(y.tolist(), [11.0])
       self.assertEqual(x.tolist(), [5.0])
 
+  def test_assign_view_corealize_order_independent(self):
+    for reader_first in (True, False):
+      x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+      reader = x + 10
+      writer = x[:1].assign(Tensor([9], dtype=dtypes.int))
+      Tensor.realize(reader, writer) if reader_first else Tensor.realize(writer, reader)
+      self.assertEqual(reader.tolist(), [11, 12])
+      self.assertEqual(x.tolist(), [9, 2])
+
+  def test_assign_view_reshape_alias_sees_write(self):
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    alias = x.reshape(2, 2)
+    prior = alias + 0
+    x[:1].assign(Tensor([9], dtype=dtypes.int))
+    post = alias + 0
+    Tensor.realize(prior, post)
+    self.assertEqual(prior.tolist(), [[1, 2], [3, 4]])
+    self.assertEqual(post.tolist(), [[9, 2], [3, 4]])
+    self.assertEqual(x.tolist(), [9, 2, 3, 4])
+
+  def test_assign_view_prior_reader_split_realize(self):
+    x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+    reader = x + 10
+    x[:1].assign(Tensor([9], dtype=dtypes.int)).realize()
+    self.assertEqual(reader.tolist(), [11, 12])
+    self.assertEqual(x.tolist(), [9, 2])
+
+  def test_assign_view_reader_between_assigns_split_realize(self):
+    x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+    x[:1].assign(Tensor([9], dtype=dtypes.int))
+    mid = x + 0
+    x[1:].assign(Tensor([8], dtype=dtypes.int))
+    x.realize()
+    self.assertEqual(mid.tolist(), [9, 2])
+    self.assertEqual(x.tolist(), [9, 8])
+
+  def test_assign_view_source_from_prior_reader_split_realize(self):
+    x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+    reader = x + 10
+    x[:1].assign(reader[1:] * 1).realize()
+    self.assertEqual(x.tolist(), [12, 2])
+    self.assertEqual(reader.tolist(), [11, 12])
+
+  def test_assign_view_source_is_prior_reader(self):
+    x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+    reader = x + 10
+    x[:].assign(reader).realize()
+    self.assertEqual(x.tolist(), [11, 12])
+    self.assertEqual(reader.tolist(), [11, 12])
+
+  def test_assign_view_prior_reader_jit_replay(self):
+    @TinyJit
+    def f(x:Tensor):
+      reader = x + 10
+      x[:1].assign(Tensor([9], device=x.device, dtype=dtypes.int))
+      return reader.realize()
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertEqual(f(x).tolist(), [i+11, i+12])
+      self.assertEqual(x.tolist(), [9, i+2])
+
+  def test_assign_view_prior_reader_is_jit_rhs(self):
+    @TinyJit
+    def f(x:Tensor):
+      reader = x[:1] + 10
+      x[:1].assign(reader)
+      return reader
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertEqual(f(x).tolist(), [i+11])
+      self.assertEqual(x.tolist(), [i+11, i+2])
+
+  def test_assign_view_realized_prior_reader_jit_replay(self):
+    @TinyJit
+    def f(x:Tensor):
+      reader = (x + 10).realize()
+      x[:1].assign(Tensor([9], device=x.device, dtype=dtypes.int))
+      return reader
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertEqual(f(x).tolist(), [i+11, i+12])
+      self.assertEqual(x.tolist(), [9, i+2])
+
+  def test_assign_view_realized_disjoint_reader_jit_replay(self):
+    @TinyJit
+    def f(x:Tensor):
+      reader = (x[:1] + 10).realize()
+      x[-1:].assign(Tensor([9], device=x.device, dtype=dtypes.int))
+      return reader
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertEqual(f(x).tolist(), [i+11])
+      self.assertEqual(x.tolist(), [i+1, 9])
+
+  def test_assign_view_explicit_effect_jit_replay(self):
+    @TinyJit
+    def f(x:Tensor):
+      x[:1].assign(Tensor([9], device=x.device, dtype=dtypes.int)).realize()
+
+    for i in range(4):
+      x = Tensor([i+1, i+2], dtype=dtypes.int).contiguous().realize()
+      self.assertIsNone(f(x))
+      self.assertEqual(x.tolist(), [9, i+2])
+
+  def test_assign_view_disjoint_reader_stays_lazy(self):
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    reader = x[:1] + 10
+    GlobalCounters.reset()
+    x[-1:].assign(Tensor([9], dtype=dtypes.int))
+    assert_kernel_count(0)
+    self.assertEqual(reader.tolist(), [11])
+    self.assertEqual(x.tolist(), [1, 2, 3, 9])
+
+  def test_assign_view_overlapping_reader_stays_lazy(self):
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    reader = x + 10
+    value = Tensor([9], dtype=dtypes.int)
+    GlobalCounters.reset()
+    x[:1].assign(value)
+    assert_kernel_count(0)
+    self.assertEqual(reader.tolist(), [11, 12, 13, 14])
+    self.assertEqual(x.tolist(), [9, 2, 3, 4])
+
+  def _count_storage_range(self, fxn):
+    # _storage_range runs a full graph_rewrite (via UOp.contiguous_view), so it dominates assign cost
+    import tinygrad.tensor as tensor_module
+    calls, orig = [], tensor_module._storage_range
+    def counting(u, base):
+      calls.append(u)
+      return orig(u, base)
+    tensor_module._storage_range = counting
+    try: fxn()
+    finally: tensor_module._storage_range = orig
+    return len(calls)
+
+  def test_assign_view_no_readers_skips_write_range(self):
+    # the write range only decides which prior readers overlap the write. with no prior readers there is
+    # nothing to decide, and computing it anyway made chained view assigns quadratic in chain length
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    n = self._count_storage_range(lambda: x[:1].assign(Tensor([9], dtype=dtypes.int)))
+    self.assertEqual(n, 0, "view assign with no prior readers must not compute the write range")
+    self.assertEqual(x.tolist(), [9, 2, 3, 4])
+
+  def test_assign_view_with_reader_still_checks_overlap(self):
+    # guard the other way: skipping the write range when readers DO exist would lose overlap detection
+    x = Tensor([1, 2, 3, 4], dtype=dtypes.int).contiguous().realize()
+    reader = x[:1] + 10
+    n = self._count_storage_range(lambda: x[-1:].assign(Tensor([9], dtype=dtypes.int)))
+    self.assertGreater(n, 0, "with a prior reader the write range is needed to classify overlap")
+    self.assertEqual(reader.tolist(), [11])
+    self.assertEqual(x.tolist(), [1, 2, 3, 9])
+
+  def test_assign_view_swap_regions(self):
+    for reverse in (False, True):
+      x = Tensor([1, 2], dtype=dtypes.int).contiguous().realize()
+      a, b = x[:1] * 1, x[1:] * 1
+      w1, w2 = x[:1].assign(b), x[1:].assign(a)
+      Tensor.realize(w2, w1) if reverse else Tensor.realize(w1, w2)
+      self.assertEqual(x.tolist(), [2, 1])
+
+  def test_assign_view_backward_prior_loss(self):
+    w = Tensor([1.0, 2.0]).contiguous().realize()
+    loss = (w * w).sum()
+    w[:1].assign(Tensor([5.0]))
+    loss.backward()
+    self.assertEqual(loss.item(), 5.0)
+    self.assertEqual(w.grad.tolist() if w.grad is not None else None, [2.0, 4.0])
+
+  def test_assign_full_view_backward_prior_loss(self):
+    for realize_first in (False, True):
+      w = Tensor([1.0, 2.0]).contiguous().realize()
+      loss = (w * w).sum()
+      w[:].assign(Tensor([5.0, 6.0]))
+      if realize_first: w.realize()
+      loss.backward()
+      if not realize_first: w.realize()
+      self.assertEqual(loss.item(), 5.0)
+      self.assertEqual(w.grad.tolist() if w.grad is not None else None, [2.0, 4.0])
+
+  def test_assign_view_prior_loss_owner_replaced(self):
+    w = Tensor([1.0, 2.0]).contiguous().realize()
+    loss = (w * w).sum()
+    w[:1].assign(Tensor([5.0]))
+    w.replace(Tensor([100.0, 200.0]).realize())
+    loss.backward()
+    self.assertIsNone(w.grad)
+
+  def test_assign_view_prior_loss_owner_replaced_with_dependency(self):
+    w = Tensor([1.0, 2.0]).contiguous().realize()
+    loss = (w * w).sum()
+    w[:1].assign(Tensor([5.0]))
+    w.replace(w + 100)
+    loss.backward()
+    self.assertIsNone(w.grad)
+
+  def test_assign_bitcast_view_realize_without_touching_base(self):
+    x = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    writer = x[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32))
+    writer.realize()
+    self.assertEqual(x.uop.buffer.numpy().tolist(), [4.0, 3.0, 3.0, 4.0])
+
   def test_assign_contiguous(self):
     b = Tensor.arange(16).reshape(4,4).clone().realize()
     a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1)

tinygrad/engine/jit.py

@@ -250,7 +250,9 @@ def __call__(self, *args, **kwargs) -> ReturnType:
       try:
         ret = self.fxn(*args, **kwargs)
         if len(params:=get_parameters(ret)): Tensor.realize(*params)
-      finally: capturing.clear()
+      finally:
+        capturing.clear()
+        Tensor._clear_capture_effects()
       if not len(self._linears): raise JitError("didn't JIT anything!")
       _check_no_non_tensor_return(ret)
       if DEBUG >= 1: print(f"JIT captured {len(self._linears)} linears with {len(input_buf_uops)} inputs")