IQ.Pilot Release Commit @ b6534c0

artifacts/package_sources/tinygrad/test/null/test_process_replay.py

@@ -0,0 +1,44 @@
+import unittest
+from tinygrad import Tensor, Device, Context
+from tinygrad.codegen import do_to_program
+from tinygrad.codegen.opt import Opt, OptOps
+from test.external.process_replay.process_replay import replay_to_program
+from test.helpers import replace_opts
+
+N = 16
+class TestProcessReplay(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
+    cls.renderer = Device[Device.DEFAULT].renderer
+
+  def test_replay_no_opts(self):
+    # opts=None means use default heuristic path
+    p = do_to_program(self.ast, self.renderer)
+    good, compare, _ = replay_to_program(p, self.ast, self.renderer)
+    self.assertEqual(good, compare)
+
+  def test_replay_empty_opts(self):
+    # opts=[] means explicitly apply zero opts (unoptimized)
+    ast = replace_opts(self.ast, [])
+    p = do_to_program(ast, self.renderer)
+    good, compare, _ = replay_to_program(p, ast, self.renderer)
+    self.assertEqual(good, compare)
+
+  def test_replay_with_opt(self):
+    # opts=[Opt(...)] means apply a specific opt
+    opts = [Opt(OptOps.UPCAST, 0, 4)]
+    ast = replace_opts(self.ast, opts)
+    p = do_to_program(ast, self.renderer)
+    good, compare, _ = replay_to_program(p, ast, self.renderer)
+    self.assertEqual(good, compare)
+
+  def test_beam(self):
+    with Context(BEAM=1):
+      ast = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
+    p = do_to_program(ast, self.renderer)
+    good, compare, _ = replay_to_program(p, ast, self.renderer)
+    self.assertEqual(good, compare)
+
+if __name__ == '__main__':
+  unittest.main(verbosity=2)