diff --git a/test/unit/test_symbolic.py b/test/unit/test_symbolic.py index 100e66a027..4f97a9397b 100644 --- a/test/unit/test_symbolic.py +++ b/test/unit/test_symbolic.py @@ -154,6 +154,12 @@ class TestSymbolic(unittest.TestCase): def test_mul_mul(self): self.helper_test_variable((Variable("a", 0, 5)*10)*9, 0, 5*10*9, "(a*90)") + def test_mul_lt(self): + self.helper_test_variable((Variable("a", 0, 5)*4)<13, 0, 1, "(a<4)") + self.helper_test_variable((Variable("a", 0, 5)*4)<16, 0, 1, "(a<4)") + self.helper_test_variable((Variable("a", 0, 5)*4)>11, 0, 1, "((a*-1)<-2)") + self.helper_test_variable((Variable("a", 0, 5)*4)>12, 0, 1, "((a*-1)<-3)") + def test_div_div(self): self.helper_test_variable((Variable("a", 0, 1800)//10)//9, 0, 20, "(a//90)") diff --git a/tinygrad/shape/symbolic.py b/tinygrad/shape/symbolic.py index d12a1b6c70..0f27e76664 100644 --- a/tinygrad/shape/symbolic.py +++ b/tinygrad/shape/symbolic.py @@ -195,6 +195,10 @@ class LtNode(OpNode): def substitute(self, var_vals: Dict[VariableOrNum, Node]) -> Node: return self.a.substitute(var_vals) < (self.b if isinstance(self.b, int) else self.b.substitute(var_vals)) class MulNode(OpNode): + def __lt__(self, b: Union[Node, int]): + if isinstance(b, Node) or isinstance(self.b, Node) or self.b == -1: return Node.__lt__(self, b) + sgn = 1 if self.b > 0 else -1 + return Node.__lt__(self.a*sgn, (b + abs(self.b) - 1)//abs(self.b)) def __mul__(self, b: Union[Node, int]): return self.a*(self.b*b) # two muls in one mul def __floordiv__(self, b: Union[Node, int], factoring_allowed=False): # NOTE: mod negative isn't handled right if self.b % b == 0: return self.a*(self.b//b)