github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

variation of (x%c)+(x//c)*c = x (#13135)

Browse Source 
when x is in the form of y//b, the idiv term might have combined
This commit is contained in:
chenyu
2025-11-06 18:53:28 -05:00
committed by GitHub
parent 42b34cf83d
commit bb8cf948f2
2 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
test/unit/test_uop_symbolic.py
View File

@@ -643,6 +643,10 @@ class TestSymbolic(unittest.TestCase):
self.helper_test_variable(lidx+(gidx//4)*8+2*(gidx%4), 0, 372, "(lidx+(gidx*2))")
self.helper_test_variable(lidx+2*(gidx%4)+(gidx//4)*8, 0, 372, "(lidx+(gidx*2))")
def test_div_mod_recombine_partial(self):
gidx = Variable("gidx", 0, 15)
self.helper_test_variable((gidx//2)%4+(gidx//8)*4, 0, 7, "gidx//2")
def test_div_mod_recombine_folded_mod(self):
a = Variable("a", 0, 2)
b = Variable("b", 0, 100)

4
tinygrad/uop/symbolic.py
View File

@@ -48,8 +48,10 @@ symbolic_simple = propagate_invalid + PatternMatcher([
(UPat.var("x") // 1, lambda x: x), # x//1 -> x
(UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base), # (x%y)%y = -> x%y (rewritten with base for speed)
# 4 variations of (x%c)+(x//c)*c = x TODO: add sorting to remove some variations
# variations of (x%c)+(x//c)*c = x TODO: add sorting to remove some variations
(UPat.var("x")%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"), lambda x,c: x), # (x%c)+(x//c)*c = x
((UPat.var("x")//UPat.cvar("a"))%UPat.cvar("c")+(UPat.var("x")//UPat.cvar("b"))*UPat.cvar("c"),
lambda x,a,b,c: x//a if a.arg*c.arg==b.arg else None), # ((x//a)%c)+(x//a*c)*c = x//a. Note if a = 1 it degenerates to the one above
((UPat.var("x")//UPat.cvar("c1"))*UPat.cvar("c3")+UPat.var("x")%UPat.cvar("c1")*UPat.cvar("c2"),
lambda x,c1,c2,c3: x*c2 if c1.arg*c2.arg==c3.arg else None), # (x%c1)*c2+(x//c1)*c3 = x*c2 if c1*c2==c3
((UPat.var("y")+(UPat.var("x")//UPat.cvar("c"))*UPat.cvar("c"))+UPat.var("x")%UPat.cvar("c"), lambda y,x,c: y+x),