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)),