sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

Better div nesting (#11811)

Browse Source 
* remove check

* use fold_divmod_congruence instead of simplify

* adjust tests

* shorten line
This commit is contained in:
Sieds Lykles
2025-08-24 04:17:40 +02:00
committed by GitHub
parent e652062f92
commit 952f729b07
2 changed files with 5 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_uop_symbolic.py
+3 -5
View File
@@ -483,9 +483,7 @@ class TestSymbolic(unittest.TestCase):
lidx2 = Variable("lidx2", 0, 3)
alu0 = gidx2*640+gidx1*160+(gidx0//5)*2+lidx0*320+lidx1*10
self.helper_test_variable((alu0+lidx2*2+1)//20, 0, 8192,
("((((((gidx0//5)+lidx2)//5)+lidx1)//2)+(((gidx2*32)+(gidx1*8))+(lidx0*16)))",
"(((lidx1+((lidx2+(gidx0//5))//5))//2)+((gidx2*32)+((gidx1*8)+(lidx0*16))))",
"((((gidx1*8)+(gidx2*32))+(lidx0*16))+((lidx1+((lidx2+(gidx0//5))//5))//2))"))
("((((gidx1*8)+(gidx2*32))+(lidx0*16))+(((lidx2+(gidx0//5))+(lidx1*5))//10))",))
def test_sum_div_complex2(self):
gidx0 = Variable("gidx0", 0, 7)
@@ -499,8 +497,8 @@ class TestSymbolic(unittest.TestCase):
gidx0 = Variable("gidx0", 0, 7)
lidx2 = Variable("lidx2", 0, 12)
lidx3 = Variable("lidx3", 0, 1)
self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)"))
self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, ("(((lidx2//2)+gidx0)//3)", "((gidx0+(lidx2//2))//3)"))
self.helper_test_variable((gidx0*4+lidx2*2+lidx3)//12, 0, 4, ("((lidx2+(gidx0*2))//6)"))
self.helper_test_variable((lidx2*2+gidx0*4+lidx3)//12, 0, 4, ("((lidx2+(gidx0*2))//6)"))
def test_sum_mul_distribute(self):
gidx0 = Variable("gidx0", 0, 7)
tinygrad/uop/symbolic.py
+2 -1
View File
@@ -202,7 +202,8 @@ def nest_div_by_smallest_factor(d: UOp, x: UOp, y: UOp) -> UOp|None:
# TODO: there are better ways to pick `div`, this sometimes adds extra divisions
# TODO: add same optimization for mod
div = min([y.arg]+[abs(f) for f in factors if abs(f) > 1 and (c%f)==0])
if (1 < div < c) and (newxs:=(newx:=(x//div)).simplify()) is not newx and x.vmin>=0 and newx.vmin>=0: return newxs//(c//div)
if (1 < div < c) and (newxs:=fold_divmod_congruence(newx:=(x//div), x, y.const_like(div))) is not None and x.vmin>=0 and newx.vmin>=0:
return newxs//(c//div)
return None
def simplify_remainder(d: UOp, x: UOp, y: UOp) -> UOp|None: