diff_schedule tests [run_process_replay] (#5958)

* diff_schedule tests [run_process_replay]

* ok to run serial

test/external/process_replay/diff_schedule.py

@@ -1,14 +1,14 @@
 # create a diff of two schedule graphs
 import difflib #, ocdiff
 from collections import defaultdict
-from typing import DefaultDict, List, Set, Tuple
-from tinygrad.engine.schedule import ScheduleItem
+from typing import DefaultDict, Dict, List, Set, Tuple
+from tinygrad.engine.schedule import LBScheduleItem, ScheduleItem
 from tinygrad.helpers import Context, colored
 from tinygrad.lazy import LazyBuffer
 from tinygrad.ops import LazyOp
-from tinygrad.engine.realize import lower_schedule_item
+from tinygrad.engine.realize import CompiledRunner, lower_schedule_item
 
-def diff_schedule(s):
+def diff_schedule(s:List[Tuple[DefaultDict[LazyBuffer, List[LazyBuffer]], Dict[LazyBuffer, LBScheduleItem]]]) -> int:
   si_for_buf: DefaultDict[LazyBuffer, List[ScheduleItem]] = defaultdict(list)
   for _,prescheduled in s:
     for ps in prescheduled.values():
@@ -24,6 +24,7 @@ def diff_schedule(s):
     #print(ocdiff.console_diff(render(ast[0]), render(ast[1])))
     ei0 = lower_schedule_item(si[0])
     ei1 = lower_schedule_item(si[1])
+    assert isinstance(ei0.prg, CompiledRunner) and isinstance(ei1.prg, CompiledRunner)
     diff = list(difflib.unified_diff(ei0.prg.p.src.splitlines(), ei1.prg.p.src.splitlines()))
     unified_diff = "\n".join(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None) for line in diff)
     print(unified_diff)
@@ -35,3 +36,4 @@ def diff_schedule(s):
     if tm_diff > 0: print(colored(f"{tm_diff:.2f}% faster", "green"))
     else: print(colored(f"{tm_diff:,.2f}% slower", "red"))
   print(f"{changed} unique kernel{'s' if changed>1 else ''} changed")
+  return changed

test/external/process_replay/test_diff_schedule.py

@@ -0,0 +1,37 @@
+import unittest
+from test.external.process_replay.diff_schedule import diff_schedule
+from tinygrad import Tensor
+from tinygrad.helpers import Context
+from tinygrad.engine.schedule import SCHEDULES
+
+class TestDiffSchedule(unittest.TestCase):
+  def test_diff_arange(self):
+    # diff a single arange kernel
+    X = Tensor.randn(10, 10).realize()
+    idxs = Tensor([0, 2]).realize()
+    xt = X[idxs]
+    with Context(ARANGE_DIFF=1): xt.schedule()
+    self.assertEqual(len(SCHEDULES), 2)
+    changed = diff_schedule(SCHEDULES)
+    self.assertEqual(changed, 1)
+    SCHEDULES.clear()
+
+    # no diff
+    a = Tensor([1])+Tensor([2])
+    with Context(ARANGE_DIFF=1): a.schedule()
+    self.assertEqual(len(SCHEDULES), 2)
+    changed = diff_schedule(SCHEDULES)
+    self.assertEqual(changed, 0)
+    SCHEDULES.clear()
+
+    # no diff with two schedule creation calls
+    a = Tensor([1])+Tensor([2])
+    with Context(ARANGE_DIFF=1): a.schedule()
+    b = Tensor([3])+Tensor([4])
+    with Context(ARANGE_DIFF=1): b.schedule()
+    self.assertEqual(len(SCHEDULES), 4)
+    changed = diff_schedule(SCHEDULES)
+    self.assertEqual(changed, 0)
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad/engine/schedule.py

@@ -4,8 +4,8 @@ from dataclasses import dataclass, field
 from typing import Tuple, List, Dict, Optional, Set, DefaultDict, cast, get_args
 from tinygrad.ops import MetaOps, BufferOps, LazyOp, Op, ReduceOps, ConstBuffer, MemBuffer, UNSAFE_PAD_OPS, UnaryOps, reduce_st
 from tinygrad.engine.graph import log_lazybuffer, realized_lazybuffer
-from tinygrad.helpers import GRAPH, DEBUG, MULTIOUTPUT, SAVE_SCHEDULE, FUSE_CONV_BW, FUSE_ARANGE, Context, GlobalCounters, colored, prod, dedup,\
-    all_int, merge_dicts, getenv, Metadata
+from tinygrad.helpers import ARANGE_DIFF, GRAPH, DEBUG, MULTIOUTPUT, SAVE_SCHEDULE, FUSE_CONV_BW, FUSE_ARANGE, Context, \
+                             GlobalCounters, colored, prod, dedup, all_int, merge_dicts, getenv, Metadata
 from tinygrad.shape.symbolic import Variable, sint
 from tinygrad.dtype import ConstType, ImageDType, dtypes
 from tinygrad.lazy import LazyBuffer
@@ -250,7 +250,7 @@ def _get_isolated_children(r:LazyBuffer, reduce_for_op:Dict[LazyBuffer, LazyBuff
   for tr in group: _recursive_group(tr, tr.st, tr, children, realizes, reduce_for_op, descendants, cache={})
   return merge_dicts([group, {} if any(tr in group for tr in descendants) else descendants])
 
-SCHEDULES: List = []
+SCHEDULES: List[Tuple[DefaultDict[LazyBuffer, List[LazyBuffer]], Dict[LazyBuffer, LBScheduleItem]]] = []
 def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]) -> \
   Tuple[DefaultDict[LazyBuffer, List[LazyBuffer]],  # this is the graph
         DefaultDict[LazyBuffer, int],               # this is the in-degree of the graph
@@ -365,7 +365,7 @@ def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]) -> \
 
   if SAVE_SCHEDULE:
     def _save():
-      if getenv("ARANGE_DIFF"):
+      if ARANGE_DIFF:
         from test.external.process_replay.diff_schedule import diff_schedule
         return diff_schedule(SCHEDULES)
       print(f"saving {len(SCHEDULES)} schedule graphs to", fp:=getenv("SAVE_SCHEDULE_PATH", "schedule.pkl"))
@@ -378,7 +378,7 @@ def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]) -> \
 
 def create_schedule_with_vars(outs:List[LazyBuffer], seen:Optional[Set[LazyBuffer]]=None) -> Tuple[List[ScheduleItem], Dict[Variable, int]]:
   if seen is None: seen = set()
-  if getenv("ARANGE_DIFF"):
+  if ARANGE_DIFF:
     with Context(FUSE_ARANGE=0, SAVE_SCHEDULE=1): _graph_schedule(outs, set())
     with Context(FUSE_ARANGE=1, SAVE_SCHEDULE=1): graph, in_degree, prescheduled = _graph_schedule(outs, seen)
   else: graph, in_degree, prescheduled = _graph_schedule(outs, seen)

tinygrad/helpers.py

@@ -108,7 +108,7 @@ GRAPH, GRAPHPATH, SAVE_SCHEDULE, RING = ContextVar("GRAPH", 0), getenv("GRAPHPAT
 MULTIOUTPUT, PROFILE, PROFILEPATH = ContextVar("MULTIOUTPUT", 1), ContextVar("PROFILE", 0), ContextVar("PROFILEPATH", temp("tinygrad_profile.json"))
 USE_TC, TC_OPT, TRANSCENDENTAL = ContextVar("TC", 1), ContextVar("TC_OPT", 0), ContextVar("TRANSCENDENTAL", 1)
 FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 0), ContextVar("FUSE_CONV_BW", 0)
-SPLIT_REDUCEOP = ContextVar("SPLIT_REDUCEOP", 1)
+SPLIT_REDUCEOP, ARANGE_DIFF = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("ARANGE_DIFF", 0)
 
 @dataclass(frozen=True)
 class Metadata: