commaai
/
tinygrad
mirror of https://github.com/commaai/tinygrad.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

uop mod-mod simplification (#5650)

This commit is contained in:
chenyu 2024-07-23 12:33:55 -04:00 committed by GitHub
parent 1384f08cd4
commit e210c87b4a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
test/unit/test_uop_symbolic.py
View File

@ -232,7 +232,6 @@ class TestSymbolic(unittest.TestCase):
def test_mul_mod_small(self):
self.helper_test_variable((Variable("a", 0, 5)*10)%9, 0, 5, "a")
@unittest.expectedFailure
def test_mod_mod(self):
self.helper_test_variable((Variable("a", 0, 31)%12)%4, 0, 3, "(a%4)")
self.helper_test_variable(((4*Variable("a", 0, 31)) % 12) % 4, 0, 0, "0")

3
tinygrad/codegen/uopgraph.py
View File

@ -245,6 +245,9 @@ constant_folder = PatternMatcher([
# mul -> (sum) -> mod
((UOp.cvar('c0')*UOp.var('x')) % UOp.cvar('c1'), lambda x,c0,c1: x*(c0.arg%c1.arg)%c1 if c0.arg >= c1.arg > 0 else None),
(((UOp.cvar('c')*UOp.var('x'))+UOp.var('x2')) % UOp.cvar('c'), lambda x,c,x2: x2%c),
# mod mod
((UOp.var('x') % UOp.cvar('c0')) % UOp.cvar('c1'), lambda x,c0,c1: x % c0 if 0 < c0.arg < c1.arg else x % c1 if c0.arg % c1.arg == 0 else None),
(((UOp.var('x') * UOp.cvar('c0')) % UOp.cvar('c1')) % UOp.cvar('c0'), lambda x,c0,c1: x.const(0)),
# two stage mul, (x*c1)*c2 = x*(c1*c2)
((UOp.var("x") * UOp.cvar("c1")) * UOp.cvar("c2"), lambda x,c1,c2: x*x.const(exec_alu(BinaryOps.MUL, x.dtype, [c1.arg, c2.arg]))),
# -(x+y) -> -x + -y