allow sinking childless consts and fold them [pr] (tinygrad#7941)

chenyuxyz · Nov 28, 2024 · aa7e167 · aa7e167
1 parent 3ab67d4
commit aa7e167
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/test/test_schedule.py b/test/test_schedule.py
@@ -16,7 +16,7 @@
 from tinygrad.ops import UOp, Ops, graph_rewrite, track_rewrites
 from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, GlobalCounters, flatten, getenv, SPLIT_REDUCEOP, unwrap, prod, Context
 from tinygrad.codegen.kernel import Kernel, verify_ast
-from tinygrad.engine.schedule import BUF_LIMIT, ScheduleItem, create_schedule, view_right, view_left
+from tinygrad.engine.schedule import BUF_LIMIT, ScheduleItem, create_schedule, view_right, view_left, do_realize
 from tinygrad.engine.realize import CompiledRunner, get_runner, run_schedule
 from tinygrad.engine.lazy import LazyBuffer, view_supported_devices
 from extra.models.llama import precompute_freqs_cis
@@ -1930,5 +1930,21 @@ def test_partial_mask(self):
     run_schedule(sched)
     np.testing.assert_allclose(b.numpy(), np.pad(a.numpy(), ((0, 5), (0, 0)))[5:])
 
+@track_rewrites(named=True)
+def big_graph_rewrite(big_graph:UOp, realizes={}) -> UOp: return graph_rewrite(big_graph, do_realize, realizes)
+class TestBigGraph(unittest.TestCase):
+  def test_sink_childless_const(self):
+    x = UOp.const(dtypes.int, 0)
+    big_graph = big_graph_rewrite(x.sink(), realizes:={})
+    self.assertIs(big_graph, UOp(Ops.SINK, dtypes.void, (x,)))
+    self.assertEqual(len(realizes), 0)
+
+  def test_sink_childless_const_alt(self):
+    x = UOp.const(dtypes.int, 0)
+    y = UOp(Ops.VIEW, dtypes.int, (UOp(Ops.BUFFER, dtypes.int.ptr(), (), 0), UOp.const(dtypes.int, 0)), ShapeTracker.from_shape((10, 10)))
+    big_graph = big_graph_rewrite(UOp.sink(x, y), realizes:={})
+    self.assertIs(big_graph, UOp(Ops.SINK, dtypes.void, (x, y)))
+    self.assertEqual(len(realizes), 1)
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)
diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py
@@ -345,7 +345,7 @@ def fold_img_cast(ctx:Dict[UOp, UOp], xb:UOp, view:UOp, b:UOp, to_cast:UOp, **kw
 
 do_realize = PatternMatcher([
   # always realize sinked ops
-  (UPat(Ops.SINK, name="sink"), lambda ctx,sink: ctx.update((x.buf_uop, x) for x in sink.src)),
+  (UPat(Ops.SINK, name="sink"), lambda ctx,sink: ctx.update((x.buf_uop, x) for x in sink.src if is_scheduled(x))),
   # always realize meta ops
   (UPatSrc({Ops.ASSIGN, Ops.CONTIGUOUS, *GroupOp.Meta}), realize),
   # realize before expand or unsafe pad ops