diff --git a/test/unit/test_uop_symbolic.py b/test/unit/test_uop_symbolic.py index 7d38180dee..b46dad7411 100644 --- a/test/unit/test_uop_symbolic.py +++ b/test/unit/test_uop_symbolic.py @@ -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) diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index f20f6f6e00..1c4f43e0aa 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -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: