rename to test_tensor_uop + use upats for asserting [pr] (#8604)

* rename to test_tensor_uop + use upats for asserting [pr]

* fix pr

test/test_tensor_uop.py

@@ -3,9 +3,9 @@ import numpy as np
 import unittest
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.engine.realize import run_schedule
-from tinygrad.ops import Ops, UOp
+from tinygrad.ops import Ops, UOp, UPat
 
-class TestLazyBuffer(unittest.TestCase):
+class TestTensorUOp(unittest.TestCase):
   def test_fromcpu_shape_tracker(self):
     def helper(a: np.ndarray):
       print(a.shape, a.strides, a.flags.c_contiguous)
@@ -68,7 +68,7 @@ class TestLazyBuffer(unittest.TestCase):
     assert lb.const_like(1).const_arg == 1.0
     assert type(lb.const_like(1).const_arg) is float
 
-  def test_forced_realized_alu(self):
+  def test_contiguous_alu(self):
     a = Tensor.randn(2, 2).realize()
     b = Tensor.randn(2, 2).realize()
     add = (a+b).contiguous()
@@ -84,13 +84,14 @@ class TestLazyBuffer(unittest.TestCase):
     sched = empty.schedule()
     self.assertEqual(len(sched), 0)
 
+reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, src=(UPat(), UPat(), UPat(Ops.REDUCE_AXIS)))))
 class TestReduceOp(unittest.TestCase):
   def test_no_split_reduce_kernel(self):
     a = Tensor.rand(4, 4).realize()
     a = a.sum()
     sched = a.schedule()
     assert len(sched) == 1
-    self.assertIs(sched[0].ast.src[0].src[2].op, Ops.REDUCE_AXIS)
+    assert reduce_kernel.match(sched[0].ast, {})
 
   def test_split_reduce_kernel_dim0(self):
     a = Tensor.rand(256, 255).realize()
@@ -98,7 +99,7 @@ class TestReduceOp(unittest.TestCase):
     sched = a.schedule()
     assert len(sched) == 2
     for s in sched:
-      self.assertIs(s.ast.src[0].src[2].op, Ops.REDUCE_AXIS)
+      assert reduce_kernel.match(s.ast, {})
 
   def test_split_reduce_kernel_dim1(self):
     a = Tensor.rand(255, 256).realize()
@@ -106,7 +107,7 @@ class TestReduceOp(unittest.TestCase):
     sched = a.schedule()
     assert len(sched) == 2
     for s in sched:
-      self.assertIs(s.ast.src[0].src[2].op, Ops.REDUCE_AXIS)
+      assert reduce_kernel.match(s.ast, {})
 
 if __name__ == "__main__":
   unittest.main()

tinygrad/engine/schedule.py

@@ -220,9 +220,6 @@ def schedule_uop(pre:UOp, ctx:ScheduleContext) -> ScheduleItem:
   sink = graph_rewrite(graph_rewrite(pre, multioutput+view_left, store_bufs:={x.buf_uop:x.src[2] for x in pre.src}), view_right)
   # remove extra uops from SINK + substitue BUFFER with DEFINE_GLOBAL
   ast = graph_rewrite(sink, to_si+append_bufs, si_ctx:=ScheduleItemContext(ctx.var_vals))
-  # capture process replay
-  if CAPTURE_PROCESS_REPLAY:
-    with Context(PICKLE_BUFFERS=0): PROCESS_REPLAY_CAPTURE[str(pre.key)] = pickle.dumps((pre, ContextVar._cache, sink))
   # deal with ASSIGN
   assign_preloads: list[UOp] = []
   if len(ctx.assigns) != 0:
@@ -238,6 +235,9 @@ def schedule_uop(pre:UOp, ctx:ScheduleContext) -> ScheduleItem:
           if len(st.views) != 1 or (mask:=st.views[0].mask) is None or ShapeTracker.from_shape(st.shape).shrink(mask) != st.shrink(mask):
             raise RuntimeError("self operand of augmented assign must be contiguous.\nhelp: consider using .contiguous():\n"
                                +colored("   - a += a.T\n", "red")+colored("   + a += a.T.contiguous()", "green"))
+  # capture process replay
+  if CAPTURE_PROCESS_REPLAY:
+    with Context(PICKLE_BUFFERS=0): PROCESS_REPLAY_CAPTURE[str(pre.key)] = pickle.dumps((pre, ContextVar._cache, ast))
   return ScheduleItem(ast, tuple(u.buffer for u in si_ctx.bufs if u.size != 0),
                       tuple(dedup(m for x in pre.toposort if (m:=ctx.ops_metadata.get(x)) is not None)), tuple(dedup(assign_preloads)))