From de04781b36d59676d1dad8a19c0f3bc89aaa219e Mon Sep 17 00:00:00 2001 From: Raine Date: Wed, 12 Aug 2026 12:39:29 -0300 Subject: [PATCH] simplify equivalent const max (#17505) * add const max folds * add regression test * move --- test/null/test_uop_symbolic.py | 5 +++++ tinygrad/uop/symbolic.py | 2 ++ 2 files changed, 7 insertions(+) diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 575fe9a55a..d13afaa2eb 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -948,6 +948,11 @@ class TestSymbolic(unittest.TestCase): self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)") self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)") + def test_equivalent_const_max(self): + x = Variable("x", -10, 10) + self.helper_test_variable((x < 0).where(0, x), 0, 10, "x.maximum(0)") + self.helper_test_variable((0 < x).where(x, 0), 0, 10, "x.maximum(0)") + def test_where_combine(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 f3d78ebfb5..a2e82fb180 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -249,6 +249,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None), (UPat(Ops.RANGE, src=(UPat(Ops.CONST,)), name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None), # max folding + ((UPat.cvar("a") < UPat.var("b")).where(UPat.var("b"), UPat.cvar("c")), lambda a,b,c: UOp.maximum(a,b) if a.val == c.val else None), + ((UPat.var("a") < UPat.cvar("b")).where(UPat.cvar("c"), UPat.var("a")), lambda a,b,c: UOp.maximum(a,b) if b.val == c.val else None), (UPat.maximum(UPat.var("x"), UPat.var("y")), lambda x,y: x if x.vmin >= y.vmax else y if x.vmax <= y.vmin else None), # TODO: why does this rule break beautiful_mnist? #((UPat.var("x")+UPat.var("z")).maximum(UPat.var("y")+UPat.var("z")), lambda x,y,z: x.maximum(y) + z),