From c59bc2951259dad7fb96517781b5838ca2d8e0fe Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Fri, 10 Mar 2023 14:16:59 -0600 Subject: [PATCH] [Arith] Add simplification rule for `x - max(x+y, z)` (#14271) This parallels an existing simplification rule for `x - min(x,y, z)`, applying the same cancellation for `max`. --- src/arith/rewrite_simplify.cc | 2 ++ tests/python/unittest/test_arith_rewrite_simplify.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc index ce2c3e1a962e..0b646ab3205a 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/arith/rewrite_simplify.cc @@ -387,6 +387,8 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const SubNode* op) { TVM_TRY_REWRITE(matches_one_of(x - min(x + y, z), x - min(y + x, z)), max(0 - y, x - z)); TVM_TRY_REWRITE(matches_one_of(x - min(z, x + y), x - min(z, y + x)), max(x - z, 0 - y)); + TVM_TRY_REWRITE(matches_one_of(x - max(x + y, z), x - max(y + x, z)), min(0 - y, x - z)); + TVM_TRY_REWRITE(matches_one_of(x - max(z, x + y), x - max(z, y + x)), min(x - z, 0 - y)); TVM_TRY_REWRITE(min(x, y) - min(y, x), ZeroWithTypeLike(x)); TVM_TRY_REWRITE(max(x, y) - max(y, x), ZeroWithTypeLike(x)); diff --git a/tests/python/unittest/test_arith_rewrite_simplify.py b/tests/python/unittest/test_arith_rewrite_simplify.py index 9be5b55ed825..dae81e8053c0 100644 --- a/tests/python/unittest/test_arith_rewrite_simplify.py +++ b/tests/python/unittest/test_arith_rewrite_simplify.py @@ -360,6 +360,10 @@ class TestSubIndex(BaseCompare): TestCase(tvm.te.max(x, y) - tvm.te.max(y, x), 0), TestCase(tvm.te.min(x, y) - tvm.te.min(x + 10, y + 10), -10), TestCase(tvm.te.min(x + 10, y + 1) - tvm.te.min(x, y - 9), 10), + TestCase(x - tvm.te.max(x + y, 0), tvm.te.min(0 - y, x)), + TestCase(x - tvm.te.max(0, x + y), tvm.te.min(x, 0 - y)), + TestCase(x - tvm.te.min(x + y, 0), tvm.te.max(0 - y, x)), + TestCase(x - tvm.te.min(0, x + y), tvm.te.max(x, 0 - y)), # DivMod patterns # truc div TestCase(x - tdiv(x, 3) * 3, tmod(x, 3)),