[tvm] branch main updated: [Arith] Add simplification rule for `x - max(x+y, z)` (#14271)

[masahi]

This is an automated email from the ASF dual-hosted git repository.

masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new c59bc29512 [Arith] Add simplification rule for `x - max(x+y, z)` 
(#14271)
c59bc29512 is described below

commit c59bc2951259dad7fb96517781b5838ca2d8e0fe
Author: Eric Lunderberg <lunderb...@users.noreply.github.com>
AuthorDate: Fri Mar 10 14:16:59 2023 -0600

    [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 ce2c3e1a96..0b646ab320 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 9be5b55ed8..dae81e8053 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)),