From 2ac5aec66bbb8fc682ebaff7439e9736b9ce009a Mon Sep 17 00:00:00 2001
From: Louis Novy <101842021+louisnovy@users.noreply.github.com>
Date: Sun, 13 Oct 2024 16:34:47 -0700
Subject: [PATCH] Fix exponential complexity in _is_padding_okay [pr] (#7008)

* preliminary test

* missed Optional

* don't check for cache during recursion

* match style from st_fixup... may be marginally faster?

* pathological test case: strongly connected DAG

* move to test_schedule as this isn't really a fusion

* oops this shouldn't be edited

* Revert "oops this shouldn't be edited"

This reverts commit 487cb027dc5120542755446d1595ec7b76c207e8.

* Revert "move to test_schedule as this isn't really a fusion"

This reverts commit 48d8c550ce84453e6fc0306e1c6c448fe1286f79.

* move to test_schedule as this isn't really a fusion

* ok no more merge error funny business
---
 test/test_schedule.py       | 10 ++++++++++
 tinygrad/engine/schedule.py |  8 +++++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/test/test_schedule.py b/test/test_schedule.py
index 19466bd6e6..858b053850 100644
--- a/test/test_schedule.py
+++ b/test/test_schedule.py
@@ -1599,6 +1599,16 @@ class TestIndexing(unittest.TestCase):
     self.assertEqual(new_uop.st, ShapeTracker.from_shape((4,)).reshape((4, 1)))
     self.assertLess(et, 1e3)
 
+  def test_strongly_connected_DAG(self):
+    val = 1.0
+    a = Tensor(val).realize()
+    def f(a):
+      for _ in range(24): a = Tensor.stack(a, a)[0]
+      return a.item()
+    r, et = timeit(f, a)
+    self.assertEqual(r, val)
+    self.assertLess(et, 1e3)
+
   def test_no_rewrite_elementwise(self):
     bufs = [UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), (), i) for i in range(3)]
     ld1 = UOp(UOps.LOAD, dtypes.int, (bufs[1], ShapeTracker.from_shape((32, 32)).to_uop()))
diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py
index fc6a2bff83..1e2a992b9d 100644
--- a/tinygrad/engine/schedule.py
+++ b/tinygrad/engine/schedule.py
@@ -238,11 +238,13 @@ def _recurse_lb(buf:LazyBuffer, realizes:Dict[LazyBuffer, None], allbufs:Dict[La
     if x.base.realized is None: children[x.base][buf] = None
     _recurse_lb(x, realizes, allbufs, simple_pads, children, assign_targets, double_reduces)
 
-def _is_padding_okay(buf:LazyBuffer, realizes:Dict[LazyBuffer, None]) -> bool:
+def _is_padding_okay(buf:LazyBuffer, realizes:Dict[LazyBuffer, None], cache:Dict[LazyBuffer, bool]) -> bool:
+  if (n:=cache.get(buf)) is not None: return n
   if buf in realizes: return True
   # NOTE: this broke to_image_idx and coder with JIT
   if buf.op in UNSAFE_PAD_OPS: return False
-  return all(_is_padding_okay(x.base, realizes) for x in buf.srcs)
+  cache[buf] = ret = all(_is_padding_okay(x.base, realizes, cache) for x in buf.srcs)
+  return ret
 
 def _recursive_group(tr:LazyBuffer, st:ShapeTracker, r:LazyBuffer, children:DefaultDict[LazyBuffer, Dict[LazyBuffer, None]],
                      realizes:Dict[LazyBuffer, None], reduce_for_op:Dict[LazyBuffer, LazyBuffer], group:Dict[LazyBuffer, None],
@@ -292,7 +294,7 @@ def _get_output_groups(outs:List[LazyBuffer]) -> \
 
   # check if we have to realize pads
   for p in simple_pads:
-    if not _is_padding_okay(p, realizes):
+    if not _is_padding_okay(p, realizes, {}):
       realizes[p] = None
 
   # find all reduces, and pair them to a elementwise op. if they can't be cleanly paired, force realize the reduce (or a contig child)