mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-07 21:13:41 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
217 lines
8.5 KiB
Python
217 lines
8.5 KiB
Python
import unittest, itertools
|
|
from tinygrad.dtype import dtypes
|
|
from tinygrad.uop.ops import Ops, UOp, GroupOp # noqa: F401
|
|
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.weakfloat), lambda x: x.rtag())])
|
|
c1 = UOp.const(1.0)
|
|
c2 = UOp.const(1)
|
|
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).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.rtag())
|
|
|
|
def test_minimum_len(self):
|
|
matcher = PatternMatcher([
|
|
(UPat(Ops.NOOP, src=(UPat(Ops.NOOP), UPat(Ops.NOOP)), allow_any_len=True), lambda: True),
|
|
])
|
|
self.assertTrue(matcher.rewrite(UOp(Ops.NOOP, src=(UOp(Ops.NOOP), UOp(Ops.NOOP), UOp(Ops.NOOP),))))
|
|
self.assertTrue(matcher.rewrite(UOp(Ops.NOOP, src=(UOp(Ops.NOOP), UOp(Ops.NOOP)))))
|
|
self.assertIsNone(matcher.rewrite(UOp(Ops.NOOP, src=(UOp(Ops.NOOP),))))
|
|
|
|
@unittest.skip("closures aren't supported on pattern matchers")
|
|
def test_match_sz_0(self):
|
|
match_cnt = 0
|
|
def fxn(x):
|
|
nonlocal match_cnt
|
|
match_cnt += 1
|
|
assert len(x.src) == 0
|
|
return UOp(Ops.CONST, src=(UOp(Ops.CONST),))
|
|
matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
|
|
c1 = UOp(Ops.CONST, arg=1.0)
|
|
# second rewrite shouldn't match anything
|
|
c1 = matcher.rewrite(c1)
|
|
c1 = matcher.rewrite(c1)
|
|
self.assertEqual(match_cnt, 1)
|
|
|
|
def test_match_sz_0_ctx(self):
|
|
def fxn(ctx, x):
|
|
ctx.append(True)
|
|
assert len(x.src) == 0
|
|
return x.replace(src=(UOp(Ops.NOOP),))
|
|
matcher = PatternMatcher([(UPat(Ops.NOOP, src=(), name="x"), fxn)])
|
|
c1 = UOp(Ops.NOOP)
|
|
# second rewrite shouldn't match anything
|
|
ctx = []
|
|
c1 = matcher.rewrite(c1, ctx)
|
|
c1 = matcher.rewrite(c1, ctx)
|
|
self.assertEqual(len(ctx), 1)
|
|
|
|
def test_uop(self):
|
|
matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())])
|
|
c1 = UOp.const(1.0)
|
|
c2 = UOp(Ops.ADD, src=(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.rtag())])
|
|
c1 = UOp.const(False)
|
|
c2 = UOp(Ops.CAST, arg=dtypes.int, src=(c1,))
|
|
c3 = UOp.const(1.0)
|
|
c4 = UOp(Ops.ADD, src=(c3, c3))
|
|
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.rtag()),
|
|
(UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()),
|
|
(UPat(Ops.MAX, name="x"), lambda x: x.rtag()),
|
|
])
|
|
c1 = UOp.const(0.0)
|
|
c2 = UOp.const(False)
|
|
c3 = UOp(Ops.MAX, src=(c1, c1))
|
|
c4 = UOp(Ops.MUL, src=(c1, c1))
|
|
c5 = UOp.const(-1)
|
|
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.rtag() if c.val in {1, -1} else None)
|
|
])
|
|
y1 = UOp.const(1)
|
|
y2 = UOp.const(2)
|
|
y3 = UOp.const(-1)
|
|
c1 = UOp(Ops.MUL, src=(y1, y2))
|
|
c2 = UOp(Ops.MUL, src=(y2, y2))
|
|
c3 = UOp(Ops.MUL, src=(y3, y2))
|
|
c4 = UOp(Ops.MUL, src=(y2, y1))
|
|
c5 = UOp(Ops.MUL, src=(y2, y3))
|
|
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
|
self.assertEqual(matcher.rewrite(c2), None)
|
|
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.rtag())])
|
|
y1 = UOp.const(1.0)
|
|
y2 = UOp.const(1.0)
|
|
c1 = UOp(Ops.ADD, src=(y1, y1))
|
|
c2 = UOp(Ops.ADD, src=(y1, y2))
|
|
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
|
self.assertEqual(matcher.rewrite(c2), c1.rtag())
|
|
|
|
def test_dtype(self):
|
|
# a concrete const dtype lives on the pair's CAST
|
|
matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
|
|
c1 = UOp.const(1.0).cast(dtypes.float32)
|
|
c2 = UOp.const(1.0).cast(dtypes.float64)
|
|
self.assertEqual(matcher.rewrite(c1), c1.rtag())
|
|
self.assertEqual(matcher.rewrite(c2), None)
|
|
|
|
def test_dtype_set(self):
|
|
matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
|
|
c1 = UOp.const(1.0).cast(dtypes.float32)
|
|
c2 = UOp.const(1.0).cast(dtypes.float64)
|
|
c3 = UOp.const(1.0).cast(dtypes.float16)
|
|
c4 = UOp.const(1).cast(dtypes.int)
|
|
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.rtag())])
|
|
c1 = UOp.const(1.0)
|
|
c2 = UOp.const(2.0)
|
|
c3 = UOp(Ops.ADD, src=(c1,c2))
|
|
self.assertEqual(matcher.rewrite(c3), c3.rtag())
|
|
self.assertEqual(matcher.rewrite(c2), None)
|
|
# that CONST/ALU -> ALU/CONST rewrite is now instant
|
|
"""
|
|
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(GroupOp.ALU))), lambda x: x)])
|
|
c4 = UOp(Ops.ADD, src=(c1,c3))
|
|
c5 = UOp(Ops.ADD, src=(c3,c1))
|
|
self.assertEqual(matcher.rewrite(c3), None)
|
|
self.assertEqual(matcher.rewrite(c4), c4)
|
|
self.assertEqual(matcher.rewrite(c5), None)
|
|
"""
|
|
|
|
def test_src_permutations(self):
|
|
matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())])
|
|
c1 = UOp.const(1.0)
|
|
c2 = UOp.const(2.0)
|
|
c3 = UOp(Ops.ADD, src=(c1,c2))
|
|
c4 = UOp(Ops.ADD, src=(c3,c2))
|
|
c5 = UOp(Ops.ADD, src=(c2,c3))
|
|
c6 = UOp(Ops.ADD, src=(c3,c4))
|
|
self.assertEqual(matcher.rewrite(c3), None)
|
|
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.rtag())])
|
|
c1 = UOp.const(1.0)
|
|
c2 = UOp.const(2.0)
|
|
c3 = UOp(Ops.ADD, src=(c1,c2))
|
|
c4 = UOp(Ops.ADD, src=(c2,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.rtag())])
|
|
c1 = UOp.const(1.0)
|
|
c2 = UOp.const(2.0)
|
|
c3 = UOp.const(3.0)
|
|
c4 = UOp(Ops.EXP2, src=(c1,))
|
|
c5 = UOp(Ops.ADD, src=(c1,c2))
|
|
c6 = UOp(Ops.MULACC, src=(c1,c2,c3))
|
|
self.assertEqual(matcher.rewrite(c4), None)
|
|
self.assertEqual(matcher.rewrite(c5), None)
|
|
self.assertEqual(matcher.rewrite(c6), c6.rtag())
|
|
|
|
def test_deep_src_permutations(self):
|
|
c1 = UOp.const(1.0)
|
|
c2 = UOp.const(2.0)
|
|
u1 = (c1 + c2) + c1
|
|
u2 = (c2 + c1) + c1
|
|
matcher = PatternMatcher([
|
|
(UPat(GroupOp.ALU, src=[UPat(GroupOp.ALU, src=[UPat(name='a'), UPat(name='b')]), UPat(name='b')]), lambda a,b: b)
|
|
])
|
|
self.assertIsNotNone(matcher.rewrite(u1))
|
|
self.assertIsNotNone(matcher.rewrite(u2))
|
|
|
|
def _assert_eq_upat(self, a:UPat, b:UPat):
|
|
assert (sorted(map(str,a.op)) if a.op else [] == (sorted(map(str,b.op)) if b.op else []))
|
|
assert (sorted(a.match_dtype) if a.match_dtype else [] == (sorted(b.match_dtype) if b.match_dtype else []))
|
|
assert (a.name, type(a.src)) == (b.name, type(b.src))
|
|
def simple_src(u:UPat):
|
|
if u.src is None: return []
|
|
if isinstance(u.src, itertools.repeat): return next(u.src[0])
|
|
return u.src[0]
|
|
for a,b in zip(simple_src(a), simple_src(b)): self._assert_eq_upat(a, b)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main(verbosity=2)
|