print_tree UPat +fix (#5132)

* fix and extend print_tree

* typing

* typing

* fix upat

* fix none

* ws

* rm prefix

* mv luop dag

* typo

* test print_tree

test/test_print_tree.py

@@ -0,0 +1,66 @@
+#%%
+import unittest
+from tinygrad.engine.graph import print_tree
+
+from tinygrad import Tensor, dtypes
+from tinygrad.codegen.uops import UOps, UOp, UPat
+from tinygrad.ops import BinaryOps
+
+import sys, io
+
+class TestPrintTree(unittest.TestCase):
+
+  def _capture_print(self, fn):
+    capturedOutput = io.StringIO()
+    sys.stdout = capturedOutput
+    fn()
+    sys.stdout = sys.__stdout__
+    return capturedOutput.getvalue()
+
+  def test_print_uop(self):
+    x = Tensor.arange(10).schedule()[-1].ast[0]
+    output =  self._capture_print(lambda: print_tree(x))
+    assert output == '\
+  0 ━┳ BufferOps.STORE MemBuffer(idx=0, dtype=dtypes.int, \
+st=ShapeTracker(views=(View(shape=(10, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)))\n\
+  1  ┗━┳ BinaryOps.ADD None\n\
+  2    ┣━┳ ReduceOps.SUM (1,)\n\
+  3    ┃ ┗━━ BufferOps.CONST ConstBuffer(val=1, dtype=dtypes.int, st=ShapeTrac\
+ker(views=(View(shape=(11, 19), strides=(0, 0), offset=0, mask=((0, 11), (9, 19))\
+, contiguous=False), View(shape=(10, 10), strides=(1, 20), offset=0, mask=None, contiguous=False))))\n\
+  4    ┗━━ BufferOps.CONST ConstBuffer(val=-1, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(10,\
+ 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)))\n'
+
+    x = UOp.var("x", dtypes.int)
+    x = (x + x) - UOp.const(dtypes.int, 2)
+    output = self._capture_print(lambda: print_tree(x))
+    assert output == '\
+  0 ━┳ UOps.ALU BinaryOps.ADD\n\
+  1  ┣━┳ UOps.ALU BinaryOps.ADD\n\
+  2  ┃ ┣━━ UOps.VAR x\n\
+  3  ┃ ┗━━ UOps.VAR x\n\
+  4  ┗━┳ UOps.ALU UnaryOps.NEG\n\
+  5    ┗━━ UOps.CONST 2\n'
+
+    x = UPat(UOp.alu(BinaryOps.ADD, UOp.var("x", dtypes.int), UOp.var("x", dtypes.int)))
+    assert self._capture_print(lambda: print_tree(x)) == '\
+  0 ━━ UOps.ALU            : dtypes.int                [<UOps.VAR: 2>, <UOps.VAR: 2>]   BinaryOps.ADD None\n'
+
+    x = UPat.compile(UOp.store(UOp.var("buf"), UOp.var("idx"),
+                               UOp(UOps.CAST, src=tuple(UOp(UOps.GEP, arg=i, src=(UOp.var("val"),)) for i in range(4)))), UOp.store)
+    assert self._capture_print(lambda: print_tree(x)) == '\
+  0 ━┳ UOps.STORE None\n\
+  1  ┣━━ None None\n\
+  2  ┣━━ None None\n\
+  3  ┗━┳ UOps.CAST None\n\
+  4    ┣━┳ UOps.GEP 0\n\
+  5    ┃ ┗━━ None None\n\
+  6    ┣━┳ UOps.GEP 1\n\
+  7    ┃ ┗━━ None None\n\
+  8    ┣━┳ UOps.GEP 2\n\
+  9    ┃ ┗━━ None None\n\
+ 10    ┗━┳ UOps.GEP 3\n\
+ 11      ┗━━ None None\n'
+
+if __name__ == "__main__":
+  unittest.main()