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

UOps pattern (x%c)+(x//c)*c = x (#6051) 

pretty cool that this is very easy to write now
This commit is contained in:
chenyu 2024-08-12 14:58:48 -04:00 committed by GitHub
parent 71c5901fc1
commit 6ed9711898
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
test/unit/test_uop_symbolic.py
View File

@ -395,6 +395,11 @@ class TestSymbolic(unittest.TestCase):
self.helper_test_variable((((alu2+150)//(-32))+4), 0, 0, "0")
self.helper_test_variable((((alu2+158)//(-32))+4), 0, 0, "0")
def test_div_mod_recombine(self):
gidx = Variable("gidx", 0, 124)
self.helper_test_variable(gidx%4+(gidx//4)*4, 0, 124, "gidx")
self.helper_test_variable((gidx//4)*4+gidx%4, 0, 124, "gidx")
@unittest.skip("not supported on uops yet")
class TestSymbolicNumeric(unittest.TestCase):
def helper_test_numeric(self, f):

2
tinygrad/codegen/uopgraph.py
View File

@ -296,6 +296,8 @@ constant_folder = PatternMatcher([
((NOp.cvar('c0')*NOp.var('x')) % NOp.cvar('c1'), lambda x,c0,c1: (x%(c1.arg//c0.arg))*c0 if c1.arg%c0.arg == 0 else None),
# mod mod
((NOp.var('x') % NOp.cvar('c0')) % NOp.cvar('c1'), lambda x,c0,c1: x % c1 if c0.arg % c1.arg == 0 else None),
# (x%c)+(x//c)*c = x
(NOp.var('x')%NOp.cvar('c')+(NOp.var('x')//NOp.cvar('c'))*NOp.cvar('c'), lambda x,c: x),
# ** combine terms **
# -(x+y) -> -x + -y
(-(NOp.var("x") + NOp.var("y")), lambda x,y: (-x)+(-y)),