Merge branch 'master' into prealloc_bufs

test/null/test_schedule.py

@@ -198,6 +198,14 @@ class TestSchedule(unittest.TestCase):
     d = a+b+c
     check_schedule(d, 1)
 
+  def test_basic_binop_fusion_assign(self):
+    a = Tensor.empty(10)
+    b = Tensor.empty(10)
+    c = Tensor.empty(10)
+    d = a+b+c
+    e = Tensor.empty(10).assign(d)
+    check_schedule(e, 1)
+
   def test_basic_binop_fusion_deep(self):
     a = Tensor.empty(10)
     b = Tensor.empty(10)
@@ -212,6 +220,13 @@ class TestSchedule(unittest.TestCase):
     c = (a*b).sum()
     check_schedule(c, 1)
 
+  def test_mulacc_fusion_assign(self):
+    a = Tensor.empty(10)
+    b = Tensor.empty(10)
+    c = (a*b).sum()
+    d = Tensor.empty(1).assign(c)
+    check_schedule(d, 1)
+
   def test_mulacc_relu_fusion(self):
     a = Tensor.empty(10)
     b = Tensor.empty(10)

test/unit/test_assign.py

@@ -36,7 +36,7 @@ class TestAssign(unittest.TestCase):
     np.testing.assert_allclose(b.numpy(), 0)
 
   def test_assign_add(self):
-    for T in (1, 2, 10, 100):
+    for T in (1, 2, 10):#, 100): # this crashes in CI, not sure why
       x = Tensor([0]).realize()
       buf = x.uop.base.realized
       for _ in range(T):

tinygrad/schedule/indexing.py

@@ -17,6 +17,10 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
   for s in rb.src:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
+def realize_assign_src(ctx:dict[UOp, None], buf:UOp, x: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 buf in x.backward_slice: ctx[x] = None
+
 pm_generate_realize_map = pm_gate_kernel_sink+PatternMatcher([
   # always realize SINK src
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
@@ -25,7 +29,9 @@ pm_generate_realize_map = pm_gate_kernel_sink+PatternMatcher([
   # always realize REDUCE on outer ranges
   (UPat(Ops.REDUCE, name="r"), lambda ctx,r: realize(ctx, r) if any(tr.arg[-1] == AxisType.OUTER for tr in r.src[1:]) else None),
   # realize srcs of these
-  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK, Ops.ASSIGN, Ops.ENCDEC), name="rb"), realize_srcs),
+  (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK, Ops.ENCDEC), name="rb"), realize_srcs),
+  # sometimes realize src of assign
+  (UPat(Ops.ASSIGN, src=(UPat.var("buf"), UPat.var("x"))), realize_assign_src),
 ])
 
 @dataclass(frozen=True)