From 46e8ea15c1e6f122fa969f8714b1e64f943b93ba Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Mon, 6 Oct 2025 09:35:50 +0800
Subject: [PATCH] split pm_substitute_recurse (#12460)

---
 test/test_schedule.py         | 1 -
 tinygrad/schedule/rangeify.py | 4 +++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/test/test_schedule.py b/test/test_schedule.py
index 7abcd13986..34a982fb99 100644
--- a/test/test_schedule.py
+++ b/test/test_schedule.py
@@ -1925,7 +1925,6 @@ class TestSchedule(unittest.TestCase):
     run_schedule(check_schedule(loss, 4))
     np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6)
 
-  @expect_rangeify_fails
   def test_const_folding_alt(self):
     t = Tensor.full((2,), 1.)
     lt = (t < 0.)
diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py
index e1c7343955..e2a7ebcb8d 100644
--- a/tinygrad/schedule/rangeify.py
+++ b/tinygrad/schedule/rangeify.py
@@ -749,7 +749,9 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   tsink = graph_rewrite(tsink, pm_rangeify, ctx=(rangeify_ctx:=RangeifyContext()), bottom_up=True, name="rangeify")
   # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right
   tsink = graph_rewrite(tsink, symbolic_simple+pm_reduce_unparented, name="symbolic")  # this supports const folding
-  tsink = graph_rewrite(tsink, pm_cleanups+pm_substitute_recurse, bottom_up=True, name="remove costly buffers")
+  tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
+  # TODO: can you substitute and remove costly buffers at the same time?
+  tsink = graph_rewrite(tsink, pm_substitute_recurse, bottom_up=True, name="run substitutes")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rangeify_ctx, name="limit buffers")
 
   # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph