cstyle strips paren for AND and OR (#10560)

test/test_renderer_failures.py

@@ -9,10 +9,11 @@ from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.runtime.ops_python import PythonRenderer
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.ops import UOp, Ops, python_alu
 from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.codegen import full_rewrite
+from tinygrad.engine.realize import lower_schedule_item
 
 def _test_uop_result(inputs:List[Tensor], stores:List[UOp], local_size=None):
   for x in inputs: x.realize()
@@ -69,6 +70,23 @@ class TestCStyleFailures(unittest.TestCase):
     ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.min(dtypes.int)+1))
     self.assertEqual(ret[0], 1)
 
+  def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
+    dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None
+    ret = Tensor.empty(1, dtype=dtype)
+    for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype))
+    schedule = ret.schedule()
+    assert len(schedule) == 1
+    ei = lower_schedule_item(schedule[0])
+    src = ei.prg.p.src
+    self.assertEqual("("*5 not in src, should_strip_paren)
+
+  def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
+  def test_repeat_mul(self): self._test_src_strip_paren(Ops.MUL)
+  def test_repeat_xor(self): self._test_src_strip_paren(Ops.XOR)
+  def test_repeat_or(self): self._test_src_strip_paren(Ops.OR)
+  def test_repeat_and(self): self._test_src_strip_paren(Ops.AND)
+  def test_repeat_sub(self): self._test_src_strip_paren(Ops.SUB, should_strip_paren=False)
+
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
 class TestWGSLFailures(unittest.TestCase):
   def test_multiply_infinity(self):

tinygrad/renderer/cstyle.py

@@ -50,8 +50,9 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.LOAD, src=(UPat.var('bidx'),), allow_any_len=True), lambda ctx,bidx: f"*{ctx[bidx]}"),
   (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var")), allow_any_len=True), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
   # alu/gep
+  # TODO: look for left-associative
   (UPat(GroupOp.ALU, name="x"), lambda ctx,x: ctx.code_for_op[x.op](
-    *([strip_parens(ctx[v]) if v.op == x.op and x.op in {Ops.ADD, Ops.MUL, Ops.XOR} else ctx[v] for v in x.src]), x.dtype)),
+    *([strip_parens(ctx[v]) if v.op == x.op and x.op in {Ops.ADD, Ops.MUL, Ops.XOR, Ops.OR, Ops.AND} else ctx[v] for v in x.src]), x.dtype)),
   (UPat(Ops.GEP, name="x"), lambda ctx,x: ctx[x.src[0]] + \
     (f"[{x.arg[0]}]" if x.src[0].dtype.count > ctx.gep_arr_threshold else f".{'xyzwabcd'[x.arg[0]]}")),
   # custom passes through with format