diff --git a/test/unit/test_uop_symbolic.py b/test/unit/test_uop_symbolic.py index 9a2fca79c3..087d848690 100644 --- a/test/unit/test_uop_symbolic.py +++ b/test/unit/test_uop_symbolic.py @@ -1019,10 +1019,7 @@ class TestSymbolicRealWorld(unittest.TestCase): #print(idx.render()) # NOTE: this used to have 13,151,129,600 in the output which is out of int32 range. self.assertIn(idx.render(), - ("((((((((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49))+(gidx0*3211264))+(gidx1*784))+(gidx2*8))+(lidx4*100352))+lidx3)+2207744)", - '((lidx3+((((((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49))+(gidx0*3211264))+(gidx1*784))+(gidx2*8))+(lidx4*100352)))+2207744)', - '((lidx3+((lidx4*100352)+((gidx2*8)+((gidx1*784)+((gidx0*3211264)+((((lidx5+1)//16)*802816)+(((lidx5+1)%16)*49)))))))+2207744)', - )) + ("(lidx3+((lidx5+1)//16*802816+(lidx5+1)%16*49+gidx0*3211264+gidx1*784+gidx2*8+lidx4*100352)+2207744)",)) class TestBounds(unittest.TestCase): def test_unrolled_arange(self): diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index ae47d756a9..a7bd23d659 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -747,9 +747,13 @@ class UOp(MathTrait, metaclass=UOpMetaClass): return fxn(**{k:v for k,v in var_vals.items() if k in varnames}) def render(self, simplify=True, pm:PatternMatcher|None=None) -> str: - with Context(TRACK_MATCH_STATS=0, SPEC=0): - ret = graph_rewrite(self.simplify() if simplify else self, renderer if pm is None else pm) - return ret.arg if ret.op is Ops.NOOP else str(ret) + ctx: dict[UOp, str] = {} + pm = renderer if pm is None else pm + for u in (s:=self.simplify() if simplify else self).toposort(): + # if there is any node in the toposort we can't render, we just render the whole thing using UOp pretty printer + if (u_str:=pm.rewrite(u, ctx=ctx)) is None: return str(s) + ctx[u] = cast(str, u_str) + return ctx[s] def pyrender(self): return pyrender(self) @@ -1249,31 +1253,36 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)]) # for debug syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<", Ops.SHR: ">>", Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"} +# comparison operators are not in here because they are chained in python, not left-associative +precedence = {Ops.NEG:0, Ops.MUL:1, Ops.IDIV:1, Ops.MOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6} +def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str: + if x.op not in precedence: return code_for_op(left, right) + return code_for_op(strip_parens(left) if precedence.get(x.src[0].op,99)<=precedence[x.op] else left, strip_parens(right) if + precedence.get(x.src[1].op,99)