rangeify: test_where_fold kernel count (#12379)

* rangeify: test_where_fold kernel count

* get these from the index

* replace ranges

* fine

* movement ops

* diff

* better

test/test_assign.py

@@ -379,8 +379,7 @@ class TestAssign(unittest.TestCase):
     a.assign(a + b)
     kc = GlobalCounters.kernel_count
     a.realize()
-    # rangeify makes two kernels
-    assert GlobalCounters.kernel_count - kc == (2 if RANGEIFY else 1)
+    assert GlobalCounters.kernel_count - kc == 1
     np.testing.assert_equal(a.numpy(), np.ones((4, 4))+np.pad(np.ones((4, 4))[:, 0:2], ((0, 0), (0, 2)), constant_values=2))
 
   def test_permuted_assignment_masked_view_not_contiguous(self):

test/test_schedule.py

@@ -1898,17 +1898,18 @@ class TestSchedule(unittest.TestCase):
       # NOTE: this is a bug on non rangeify
       np.testing.assert_equal(tst.numpy(), a.numpy())
 
-  def test_setitem_sched(self, transpose=False):
+  def test_setitem_sched(self, mop=lambda x:x, expected_kcount=1):
     a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize()
-    a2 = a.T if transpose else a
+    a2 = mop(a)
     expected = (a+a2).tolist()
     a.assign(a+a2)
     kcount = len(sched:=a.schedule())
     run_schedule(sched)
     self.assertListEqual(a.tolist(), expected)
-    self.assertEqual(kcount, 2 if transpose else 1)
+    self.assertEqual(kcount, expected_kcount)
   @unittest.skipUnless(RANGEIFY>0, "this asserts on non rangeify")
-  def test_setitem_permuted_sched(self): self.test_setitem_sched(transpose=True)
+  def test_setitem_permuted_sched(self): self.test_setitem_sched(lambda x: x.T, 2)
+  def test_setitem_paddded_sched(self): self.test_setitem_sched(lambda x: x.shrink_to(4, 1).pad_to(4, 4), 1)
 
   def test_sparse_categorical_crossentropy_simple(self):
     X = Tensor([[0, 2, 3], [1, 2, 3]]).realize()

tinygrad/schedule/rangeify.py

@@ -16,6 +16,13 @@ ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER,
                      Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL,
                      Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL}
 
+def find_permutes(a:UOp, b:UOp, assign:UOp):
+  if not (permutes:=[s for s in b.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS)
+                     if s.op in GroupOp.Movement and s.op not in {Ops.RESHAPE, Ops.EXPAND, Ops.PAD, Ops.SHRINK}]): return
+  target = a.base
+  for p in permutes:
+    if any(s is target for s in p.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS-{Ops.BUFFER})): return assign.replace(src=(a, b.contiguous()))
+
 earliest_rewrites = PatternMatcher([
   # just removing it works...
   (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]),
@@ -46,9 +53,8 @@ earliest_rewrites = PatternMatcher([
    lambda x,target,assign: x.f(Ops.CONTIGUOUS, tag=assign.tag) if ((t:=target.base).op is not Ops.BUFFER and \
        not (t.op is Ops.MSTACK and all(s.op is Ops.BUFFER for s in t.src))) else None),
 
-  # realize before assign if input permutes the target buffer
-  (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), lambda a,b,assign: assign.replace(src=(a, b.contiguous())) \
-      if any(x.base is a.base and x is not a for x in b.toposort(gate=lambda x:x.op not in ALWAYS_CONTIGUOUS)) else None),
+   # realize before assign if input permutes the target buffer
+   (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), find_permutes),
 
   # copy only to different device
   (UPat(Ops.COPY, src=(UPat.var("x"), UPat()), name="copy"), lambda x,copy: x.f(Ops.NOOP, tag=copy.tag) if x.device == copy.device else None),