[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`.
apache · Mar 10, 2023 · c59bc29 · c59bc29
1 parent 852f97d
commit c59bc29
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git 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
@@ -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)),