diff --git a/test/test_renderer_failures.py b/test/test_renderer_failures.py index 329975b94e..d04aecd6d4 100644 --- a/test/test_renderer_failures.py +++ b/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): diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index debdd979c6..eae100845a 100644 --- a/tinygrad/renderer/cstyle.py +++ b/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