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

MulNode.__lt__ rule (#2086)

Browse Source 
* Added the rule

* Added tests

* flake8

* self.b == -1 shortcut
This commit is contained in:
Umut Zengin
2023-10-17 13:18:35 -07:00
committed by GitHub
parent f76fbd23e9
commit 01b98b7f42
2 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_symbolic.py
+6
View File
@@ -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)")
tinygrad/shape/symbolic.py
+4
View File
@@ -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)