From 17557d7fdfe43c3bbcb5e1b51621733cb0089b54 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 22 Jul 2026 11:08:14 -0400 Subject: [PATCH] bring back fold_where_closure [pr] (#17134) * bring back fold_where_closure [pr] generic enough to delete delete pm_index_invalid * cannot cache self --- test/null/test_const_folding.py | 2 +- test/null/test_simplify_valid_idx.py | 4 +-- test/null/test_uop_symbolic.py | 48 +++++++++++++++++++++------- tinygrad/uop/ops.py | 6 ++++ tinygrad/uop/symbolic.py | 20 ++++++------ 5 files changed, 55 insertions(+), 25 deletions(-) diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index 2f296c0998..c15b38fdae 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -50,7 +50,7 @@ class TestWeakConstFolding(unittest.TestCase): self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75)) def test_invalid_poison(self): - self.assertIs(UOp.const(dtypes.weakint, Invalid).alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid) + self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid) class TestBinaryOpsConstFolding(unittest.TestCase): def test_add_literal_zero(self): diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index a67e5fb49e..6ebf4b2991 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -568,7 +568,7 @@ class TestRangeShrink(unittest.TestCase): from tinygrad.dtype import Invalid r = Range(0, 204) x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid) - ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, 0)).sink()) + ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 4) @@ -577,7 +577,7 @@ class TestRangeShrink(unittest.TestCase): from tinygrad.dtype import Invalid r = Range(0, 204) x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid) - ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(0, x)).sink()) + ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 4) diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 6acf2cad65..d75d56bc06 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -1001,7 +1001,6 @@ class TestSymbolic(unittest.TestCase): # (a if ((s<5)&(s<6)) else b) -> (a if (s<5) else b) self.helper_test_variable(expr, 0, 3, "(s<5).where(a, b)") - @unittest.expectedFailure def test_where_closure_folding(self): # cond.where(t, f) where f contains cond.where(a, b) should fold the inner where to b in false branch x = Variable("x", 0, 10) @@ -1011,6 +1010,41 @@ class TestSymbolic(unittest.TestCase): # the inner where should be folded: true branch gets -x, false branch gets x self.helper_test_variable(outer, -20, 11, "(x<5).where((x*-2), (x+1))") + def test_where_closure_folding_deep(self): + x = Variable("x", 0, 10) + cond = x < 5 + w1 = cond.where(-x, x) + w2 = cond.where(w1*2, w1+1) + self.helper_test_variable(cond.where(w2*3, w2+7), -60, 18, "(x<5).where((x*-6), (x+8))") + + def test_where_closure_folding_different_cond(self): + # a nested where on a different condition is not folded + x = Variable("x", 0, 10) + a = Variable("a", 0, 3) + b = Variable("b", 0, 3) + expr = (x<5).where((x<7).where(a, b), (x<7).where(b, a)) + self.helper_test_variable(expr, 0, 3, "(x<5).where((x<7).where(a, b), (x<7).where(b, a))") + + def test_where_closure_folding_derived_cond(self): + # cond is a value inside the branch: (!cond).where(a, b) is b in the true branch + x = Variable("x", 0, 10) + a = Variable("a", 0, 3) + b = Variable("b", 0, 3) + c = Variable("c", 0, 3) + expr = (x<5).where((x<5).logical_not().where(a, b)*2, c) + self.helper_test_variable(expr, 0, 6, "(x<5).where((b*2), c)") + + def test_where_closure_folding_valid(self): + # a valid gate on the same cond folds in the true branch, the false branch keeps the Invalid gate + x = Variable("x", 0, 10) + a = Variable("a", 0, 3) + cond = x < 5 + expr = cond.where(a.valid(cond), Variable("c", 0, 3)) + self.assertIs(graph_rewrite(expr, sym), cond.where(a, UOp.invalid())) + # a same-cond valid gate in the false branch is Invalid there + expr = cond.where(Variable("t", 0, 3), a.valid(cond)) + self.assertIs(graph_rewrite(expr, sym), cond.where(Variable("t", 0, 3), UOp.invalid())) + def test_symbolic_div(self): # from symbolic arange a = Variable("a", 1, 10) @@ -1259,18 +1293,10 @@ class TestInvalidIndex(unittest.TestCase): idx = (ridx<5).where(ridx, UOp.invalid())*0 self.assertIs(idx.simplify(), (ridx<5).where(0, UOp.invalid()), "multiplying an index by 0 should preserve the invalid") - def test_invalid_comparison_drops_invalid(self): - # comparisons return a bool, and bools can't be invalid - ridx = Variable("ridx", 0, 10) - idx = (ridx<5).where(ridx, UOp.invalid())<3 - self.assertIs(idx.simplify(), (ridx<3), "comparison of index should drop the invalid") - self.assertIs(idx.where(UOp.const(dtypes.int, 1), 0).simplify(), (ridx<3).where(UOp.const(dtypes.int, 1), 0), - "comparison of index should drop the invalid") - def test_alu_moves_inside_invalid(self): ridx = Variable("ridx", 0, 10) - idx = (ridx<5).where(ridx, UOp.invalid())*10 - self.assertIs(idx.simplify(), (ridx<5).where(ridx*10, UOp.invalid()), "multiplying an index by 0 should preserve the invalid") + self.assertIs((10*(ridx<5).where(ridx, UOp.invalid())).simplify(), (ridx<5).where(ridx*10, UOp.invalid()), + "Invalid should poison either binary operand position") def test_merge_invalid_conditions(self): ridx0 = Variable("ridx0", 0, 10) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 1f8ed4aa98..47de43766b 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -274,6 +274,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # Check self first, then iterate backward_slice (avoids creating intermediate dict) return self.op in ops or any(x.op in ops for x in self.backward_slice) + @recursive_property + def _bool_slice(self) -> frozenset[UOp]: return frozenset().union(*[s.bool_slice for s in self.src]) + # NOTE: self is added outside the cache, a cached self-reference is a cycle the refcounter can't free + @property + def bool_slice(self) -> frozenset[UOp]: return self._bool_slice | {self} if self.dtype is dtypes.bool else self._bool_slice + def toposort(self, gate:Callable|None=None, enter_calls=True) -> dict[UOp, None]: cache: dict[UOp, None] = {} stack: list[tuple[UOp, bool]] = [(self, False)] # each stack entry is (node, visited_flag) diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 854b30e17b..ee00b9976a 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -215,6 +215,13 @@ commutative = PatternMatcher([ x.replace(src=x.src[::-1]) if x.src[1].tuplize < x.src[0].tuplize and not x.src[0].tuplize < x.src[1].tuplize else None), ]) +def fold_where_closure(cond:UOp, t:UOp, f:UOp) -> UOp|None: + """in cond.where(t, f), cond is True within t and False within f""" + if cond not in t.bool_slice and cond not in f.bool_slice: return None + # INDEX gates are owned by the valid/store-coalescing machinery, leave them alone + if any(u.op_in_backward_slice_with_self(Ops.INDEX) for u in (cond, t, f)): return None + return cond.where(t.substitute({cond: cond.const_like(True)}), f.substitute({cond: cond.const_like(False)})) + symbolic = symbolic_simple+commutative+PatternMatcher([ # ** boolean algebra ** # TODO: make a more general or folder like simplify_valid @@ -233,6 +240,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ # ** where folding ** (UPat.var("cond", dtype=dtypes.bool).logical_not().where(UPat.var("t"), UPat.var("f")), lambda cond, t, f: cond.where(f,t) if f.arg is not Invalid else None), + # in cond.where(t, f), uses of cond fold to True within t and False within f + (UPat.var("cond", dtype=dtypes.bool).where(UPat.var("t"), UPat.var("f")), fold_where_closure), # alu of two where with same conds can combine, only do if true branch or false branch is const (UPat(GroupOp.Binary, name="alu", src=(UPat.var("c").where(UPat.var("t"), UPat.var("f")), UPat.var("c").where(UPat.var("tt"), UPat.var("ff")))), \ lambda alu,c,t,tt,f,ff: c.where(t.alu(alu.op, tt), f.alu(alu.op, ff)) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None), @@ -404,15 +413,6 @@ def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None: if IMAGE.value > 0 and x.op_in_backward_slice_with_self(Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD): return None return cond.where(uop_given_valid(cond, x, try_simplex=False), i) -# TODO: this is O(number of WHERE * number of node) -# def fold_where_closure(cond:UOp, t:UOp, f:UOp) -> UOp|None: -# """In cond.where(t, f), fold nested cond.where(a, b) -> a in t, -> b in f""" -# def is_valid_where(u:UOp) -> bool: return u.op is Ops.WHERE and u.src[0] is cond and Invalid not in (u.src[1].arg, u.src[2].arg) -# t_subs, f_subs = {u: u.src[1] for u in t.toposort() if is_valid_where(u)}, {u: u.src[2] for u in f.toposort() if is_valid_where(u)} -# if not t_subs and not f_subs: return None -# new_t, new_f = t.substitute(t_subs).simplify() if t_subs else t, f.substitute(f_subs).simplify() if f_subs else f -# return None if new_t is t and new_f is f else cond.where(new_t, new_f) - pm_simplify_valid = PatternMatcher([ # simplify valid (UPat(Ops.AND, name="valid"), simplify_valid), @@ -434,8 +434,6 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([ (UPat(GroupOp.ALU, src=(UPat(Ops.STACK, src=UPat(name='x')), UPat(Ops.STACK, src=UPat(name='y'))), name='alu'), lambda x,y,alu: UOp(Ops.STACK, src=(UOp(alu.op, src=(x,y)),))), # ** where ** - # # fold nested where with same condition: in cond.where(t,f), cond.where(a,b)->a in t, ->b in f - # (UPat.var("cond").where(UPat.var("t"), UPat.var("f")), fold_where_closure), # push cast to branches (UPat.var("s").where(UPat.var("a"), UPat.var("b")).cast().named("cast"), lambda s,a,b,cast: s.where(a.cast(cast.dtype), b.cast(cast.dtype))), # ** pow **