From 90b1c0dd9676e91376f577e8a65eeb7338f2eb12 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Wed, 1 Oct 2025 09:35:12 +0300 Subject: [PATCH] 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 | 3 +-- test/test_schedule.py | 9 +++++---- tinygrad/schedule/rangeify.py | 12 +++++++++--- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/test/test_assign.py b/test/test_assign.py index b35223fc12..63b6227b9e 100644 --- a/test/test_assign.py +++ b/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): diff --git a/test/test_schedule.py b/test/test_schedule.py index a0909a105f..d463c8aafb 100644 --- a/test/test_schedule.py +++ b/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() diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 38242849da..9af7b276ff 100644 --- a/tinygrad/schedule/rangeify.py +++ b/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),