tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 10:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

if upat returns self, it's none (#10898)

Browse Source 
* if upat returns self, it's none

* fix pm tests
This commit is contained in:
George Hotz
2025-06-20 12:11:19 -07:00
committed by GitHub
parent 4f179b9ddb
commit fc9f883870
3 changed files with 46 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_graph_rewrite.py
+7 -1
View File
@@ -1,7 +1,7 @@
import unittest, math
from tinygrad import dtypes
from tinygrad.helpers import all_same
from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu
from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, UPat
from tinygrad.codegen import full_rewrite_to_sink
# Helper function to apply the graph rewrite
@@ -284,5 +284,11 @@ class TestSubstitute(unittest.TestCase):
# the srcs are rewritten but we keep tag
self.assertIs(ret, (b+4).replace(tag=1))
class TestRecurse(unittest.TestCase):
def test_no_inf_loop(self):
a = UOp.variable('a', 0, 10)
pm = PatternMatcher([(UPat(Ops.DEFINE_VAR, name="x"), lambda x: x)])
graph_rewrite(a, pm)
if __name__ == '__main__':
unittest.main()
test/unit/test_pattern_matcher.py
+37 -37
View File
@@ -5,23 +5,23 @@ from tinygrad.uop.ops import PatternMatcher, UPat
class TestPatternMatcher(unittest.TestCase):
def test_simple_match(self):
matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x)])
matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x.rtag())])
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c2 = UOp(Ops.CONST, dtypes.int, arg=1)
self.assertEqual(matcher.rewrite(c1), c1)
self.assertEqual(matcher.rewrite(c1), c1.rtag())
self.assertEqual(matcher.rewrite(c2), None)
def test_upat_any(self):
def test(a, x=None, y=None, z=None):
#print(x,y,z)
if y is not None: return a+y
if y is not None: return (a+y).rtag()
matcher = PatternMatcher([
(UPat.var("a")+UPat.any(UPat.var("x"), UPat.var("y"), UPat.var("z")), test),
])
v1 = UOp.variable("a", 0, 10)
v2 = UOp.variable("b", 0, 10)
c1 = v1+v2
self.assertEqual(matcher.rewrite(c1), c1)
self.assertEqual(matcher.rewrite(c1), c1.rtag())
def test_minimum_len(self):
matcher = PatternMatcher([
@@ -60,43 +60,43 @@ class TestPatternMatcher(unittest.TestCase):
self.assertEqual(len(ctx), 1)
def test_uop(self):
matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x)])
matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())])
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c2 = UOp(Ops.ADD, dtypes.float, (c1, c1))
self.assertEqual(matcher.rewrite(c1), c1)
self.assertEqual(matcher.rewrite(c1), c1.rtag())
self.assertEqual(matcher.rewrite(c2), None)
def test_uop_set(self):
matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x)])
matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())])
c1 = UOp(Ops.CONST, dtypes.bool, arg=False)
c2 = UOp(Ops.CAST, dtypes.int, (c1,))
c3 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c4 = UOp(Ops.ADD, dtypes.float, (c3, c3))
self.assertEqual(matcher.rewrite(c1), c1)
self.assertEqual(matcher.rewrite(c2), c2)
self.assertEqual(matcher.rewrite(c1), c1.rtag())
self.assertEqual(matcher.rewrite(c2), c2.rtag())
self.assertEqual(matcher.rewrite(c4), None)
def test_arg(self):
matcher = PatternMatcher([
(UPat(Ops.CONST, arg=0, name="x"), lambda x: x),
(UPat(Ops.CONST, arg=False, name="x"), lambda x: x),
(UPat(Ops.MAX, name="x"), lambda x: x),
(UPat(Ops.CONST, arg=0, name="x"), lambda x: x.rtag()),
(UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()),
(UPat(Ops.MAX, name="x"), lambda x: x.rtag()),
])
c1 = UOp(Ops.CONST, dtypes.float, arg=0.0)
c2 = UOp(Ops.CONST, dtypes.bool, arg=False)
c3 = UOp(Ops.MAX, dtypes.float, (c1, c1))
c4 = UOp(Ops.MUL, dtypes.float, (c1, c1))
c5 = UOp(Ops.CONST, dtypes.int, arg=-1)
self.assertEqual(matcher.rewrite(c1), c1)
self.assertEqual(matcher.rewrite(c2), c2)
self.assertEqual(matcher.rewrite(c3), c3)
self.assertEqual(matcher.rewrite(c1), c1.rtag())
self.assertEqual(matcher.rewrite(c2), c2.rtag())
self.assertEqual(matcher.rewrite(c3), c3.rtag())
self.assertEqual(matcher.rewrite(c4), None)
self.assertEqual(matcher.rewrite(c5), None)
def test_filter_arg(self):
matcher = PatternMatcher([
(UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"),
lambda x,c: x if c.arg in {1, -1} else None)
lambda x,c: x.rtag() if c.arg in {1, -1} else None)
])
y1 = UOp(Ops.CONST, dtypes.int, arg=1)
y2 = UOp(Ops.CONST, dtypes.int, arg=2)
@@ -106,45 +106,45 @@ class TestPatternMatcher(unittest.TestCase):
c3 = UOp(Ops.MUL, dtypes.int, (y3, y2))
c4 = UOp(Ops.MUL, dtypes.int, (y2, y1))
c5 = UOp(Ops.MUL, dtypes.int, (y2, y3))
self.assertEqual(matcher.rewrite(c1), c1)
self.assertEqual(matcher.rewrite(c1), c1.rtag())
self.assertEqual(matcher.rewrite(c2), None)
self.assertEqual(matcher.rewrite(c3), c3)
self.assertEqual(matcher.rewrite(c4), c4)
self.assertEqual(matcher.rewrite(c5), c5)
self.assertEqual(matcher.rewrite(c3), c3.rtag())
self.assertEqual(matcher.rewrite(c4), c4.rtag())
self.assertEqual(matcher.rewrite(c5), c5.rtag())
def test_dup_name(self):
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x)])
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())])
y1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
y2 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c1 = UOp(Ops.ADD, dtypes.float, (y1, y1))
c2 = UOp(Ops.ADD, dtypes.float, (y1, y2))
self.assertEqual(matcher.rewrite(c1), c1)
self.assertEqual(matcher.rewrite(c2), c1)
self.assertEqual(matcher.rewrite(c1), c1.rtag())
self.assertEqual(matcher.rewrite(c2), c1.rtag())
def test_dtype(self):
matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x)])
matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c2 = UOp(Ops.CONST, dtypes.float64, arg=1.0)
self.assertEqual(matcher.rewrite(c1), c1)
self.assertEqual(matcher.rewrite(c1), c1.rtag())
self.assertEqual(matcher.rewrite(c2), None)
def test_dtype_set(self):
matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x)])
matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c2 = UOp(Ops.CONST, dtypes.float64, arg=1.0)
c3 = UOp(Ops.CONST, dtypes.float16, arg=1.0)
c4 = UOp(Ops.CONST, dtypes.int, arg=1)
self.assertEqual(matcher.rewrite(c1), c1)
self.assertEqual(matcher.rewrite(c2), c2)
self.assertEqual(matcher.rewrite(c1), c1.rtag())
self.assertEqual(matcher.rewrite(c2), c2.rtag())
self.assertEqual(matcher.rewrite(c3), None)
self.assertEqual(matcher.rewrite(c4), None)
def test_src_one(self):
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x)])
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())])
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c2 = UOp(Ops.CONST, dtypes.float, arg=2.0)
c3 = UOp(Ops.ADD, dtypes.float, (c1,c2))
self.assertEqual(matcher.rewrite(c3), c3)
self.assertEqual(matcher.rewrite(c3), c3.rtag())
self.assertEqual(matcher.rewrite(c2), None)
# that CONST/ALU -> ALU/CONST rewrite is now instant
"""
@@ -157,7 +157,7 @@ class TestPatternMatcher(unittest.TestCase):
"""
def test_src_permutations(self):
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x)])
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())])
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c2 = UOp(Ops.CONST, dtypes.float, arg=2.0)
c3 = UOp(Ops.ADD, dtypes.float, (c1,c2))
@@ -165,21 +165,21 @@ class TestPatternMatcher(unittest.TestCase):
c5 = UOp(Ops.ADD, dtypes.float, (c2,c3))
c6 = UOp(Ops.ADD, dtypes.float, (c3,c4))
self.assertEqual(matcher.rewrite(c3), None)
self.assertEqual(matcher.rewrite(c4), c4)
self.assertEqual(matcher.rewrite(c5), c5)
self.assertEqual(matcher.rewrite(c4), c4.rtag())
self.assertEqual(matcher.rewrite(c5), c5.rtag())
self.assertEqual(matcher.rewrite(c6), None)
def test_src_repeat(self):
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x)])
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())])
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c2 = UOp(Ops.CONST, dtypes.float, arg=2.0)
c3 = UOp(Ops.ADD, dtypes.float, (c1,c2))
c4 = UOp(Ops.ADD, dtypes.float, (c2,c3))
self.assertEqual(matcher.rewrite(c3), c3)
self.assertEqual(matcher.rewrite(c3), c3.rtag())
self.assertEqual(matcher.rewrite(c4), None)
def test_allow_len(self):
matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x)])
matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())])
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
c2 = UOp(Ops.CONST, dtypes.float, arg=2.0)
c3 = UOp(Ops.CONST, dtypes.float, arg=3.0)
@@ -188,7 +188,7 @@ class TestPatternMatcher(unittest.TestCase):
c6 = UOp(Ops.MULACC, dtypes.float, (c1,c2,c3))
self.assertEqual(matcher.rewrite(c4), None)
self.assertEqual(matcher.rewrite(c5), None)
self.assertEqual(matcher.rewrite(c6), c6)
self.assertEqual(matcher.rewrite(c6), c6.rtag())
def test_deep_src_permutations(self):
c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
tinygrad/uop/ops.py
+2 -1
View File
@@ -90,6 +90,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
assert len(kwargs) == 0, f"unused kwargs in replace {list(kwargs)}"
if (self.op, self.dtype, self.src, self.arg, self.tag) == new_args: return self
return UOp(*new_args)
def rtag(self, tag=True): return self.replace(tag=tag)
@functools.cached_property
def key(self) -> bytes:
return hashlib.sha256(str((self.op, self.dtype, self.arg)).encode() + b"".join([s.key for s in self.src])).digest()
@@ -708,7 +709,7 @@ class PatternMatcher:
ler = {u.op for u in uop.src}
for _,match,early_reject in self.pdict.get(uop.op, []):
if not early_reject.issubset(ler): continue
if (ret:=match(uop, ctx)) is not None: return ret
if (ret:=match(uop, ctx)) is not None and ret is not uop: return ret
return None
def fixed_point_rewrite(self, uop:UOp, ctx=None) -> UOp: