diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 48c5eec645..d1c8bb3bfc 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -967,6 +967,28 @@ class TestSymbolic(unittest.TestCase): # not combining # TODO: can combine if one is identity element const self.helper_test_variable(aa+ab, 0, 6, "((x<2).where(a, b)+(x<2).where(a, 0))") + def test_where_combine_cross_zero(self): + cond = Variable("x", 0, 3) < 2 + a = Variable("a", 0, 3) + b = Variable("b", 0, 3) + self.helper_test_variable(cond.where(a, a.ufix(0)) + cond.where(b.ufix(0), b), 0, 3, "(x<2).where(a, b)") + self.helper_test_variable(cond.where(a, a.ufix(0)) + cond.where(a.ufix(0), a), 0, 3, "a") + + def test_where_or_dual(self): + m1 = Variable("x", 0, 3) < 2 + m2 = Variable("y", 0, 3) < 2 + a = Variable("a", 0, 3) + b = Variable("b", 0, 3) + self.helper_test_variable(m1.where(a, m2.where(a, b)), 0, 3, "((x<2)|(y<2)).where(a, b)") + + def test_bool_ne_false(self): + cond = Variable("x", 0, 3) < 2 + self.helper_test_variable(cond.ne(False), 0, 1, "(x<2)") + + def test_bitcast_chain(self): + a = Variable("a", 0, 3) + self.assertIs(graph_rewrite(a.bitcast(dtypes.float32).bitcast(a.dtype), sym), a) + def test_negation_in_where(self): cond = Variable("x", 0, 3) < 2 a = Variable("a", 0, 3) diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 1106c5a3a3..8b72c88622 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -109,6 +109,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat(Ops.ADD, dtype=dtypes.weakint, name="x"), fold_add_divmod_recombine), (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c"), lambda x,c: x if c.arg else c), (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c"), lambda x,c: c if c.arg else x), + (UPat.var("x", dtype=dtypes.bool) != UPat.const(dtypes.bool, False), lambda x: x), # x != False -> x (UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x), (UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x), (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x), @@ -156,6 +157,8 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast), # b.cast(a).cast(b) -> b if a preserves all values in b (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_lossless_cast(b.dtype, a.dtype) else None), + # bitcast twice + (UPat(Ops.BITCAST, name="b", src=(UPat.var('x').bitcast(),)), lambda x,b: x.bitcast(b.dtype)), (UPat.var("x").cast(dtypes.bool), lambda x: x != 0), # ** pow ** (UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow), @@ -174,6 +177,8 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1), # a.where(b.where(c, d), d) -> (a & b).where(c, d) (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)), + # a.where(c, b.where(c, d)) -> (a | b).where(c, d) + (UPat.var("a").where(UPat.var("c"), UPat.var("b").where(UPat.var("c"), UPat.var("d"))), lambda a,b,c,d: (a|b).where(c,d)), ])+mop_cleanup # ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ******** @@ -237,6 +242,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ # if its a plus we add the associative variation too ((UPat.var("y")+UPat.var("c").where(UPat.var("t"), UPat.var("f"))) + UPat.var("c").where(UPat.var("tt"), UPat.var("ff")), \ lambda y,c,t,tt,f,ff: y+c.where(t+tt, f+ff) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None), + # complementary zero branches under the same condition select directly + (UPat.var("c").where(UPat.var("t"), 0) + UPat.var("c").where(0, UPat.var("f")), lambda c,t,f: c.where(t, f)), # ALU/variable min==max -> CONST (UPat({Ops.CMPLT, Ops.CMPNE, Ops.FLOORDIV, Ops.FLOORMOD, Ops.PARAM, Ops.BIND, Ops.SPECIAL}, name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),