diff --git a/test/unit/test_uop_symbolic.py b/test/unit/test_uop_symbolic.py index 367e18e0bf..b693f80da4 100644 --- a/test/unit/test_uop_symbolic.py +++ b/test/unit/test_uop_symbolic.py @@ -351,9 +351,8 @@ class TestSymbolic(unittest.TestCase): def test_sum_div_partial_remove(self): self.helper_test_variable(usum([Variable("idx0", 0, 127)*4, Variable("idx2", 0, 3)])//4, 0, 127, "idx0") - # TODO: this is wrong def test_div_numerator_negative(self): - self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -9, 0, "(idx*-1)") + self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)") def test_div_into_mod(self): self.helper_test_variable((Variable("idx", 0, 16)*4)%8//4, 0, 1, "(idx%2)") diff --git a/tinygrad/codegen/symbolic.py b/tinygrad/codegen/symbolic.py index b49ca66dd7..d97cf7dace 100644 --- a/tinygrad/codegen/symbolic.py +++ b/tinygrad/codegen/symbolic.py @@ -152,7 +152,7 @@ def div_and_mod_folding(x: UOp, y: UOp, which: Literal[Ops.MOD, Ops.IDIV], split # a//c = (a-a%c)/c, if we can fold a%c, we can fold a//c # within a mod we can freely subtract multiples of c, we use this to see if a is congruent to an expression whose vmin/vmax are between 0 and c rems = [min(r, r-c, key=abs) for r in remainders] - if (rem:=sum(r*v for r,v in zip(rems,svars))+const%c).vmin//c==rem.vmax//c: + if (rem:=sum(r*v for r,v in zip(rems,svars))+const%c).vmin//c==rem.vmax//c and all(f > 0 for f in factors): if which is Ops.MOD: return rem - rem.vmin//c*c return sum((f-r)//c * v for f,r,v in zip(factors,rems,svars)) + (const-const%c+rem.vmin//c*c)//c diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 4ee8f067c5..b74102d96b 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -598,6 +598,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): return False # False if not sure def const_factor(self) -> int: """largest known int that divides self""" + # TODO: for negatives it's not the largest if self.op is Ops.CONST: return self.arg if self.op is Ops.VCONST: return math.gcd(*self.arg) if self.op is Ops.ADD: return math.gcd(self.src[0].const_factor(), self.src[1].const_factor())