Additional canonicalization added for AddNode (#5846)

apache · Jun 18, 2020 · 24df0ba · 24df0ba
1 parent d85efa1
commit 24df0ba
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/python/tvm/tir/expr.py b/python/tvm/tir/expr.py
@@ -61,7 +61,7 @@ def __add__(self, other):
         return _generic.add(self, other)
 
     def __radd__(self, other):
-        return self.__add__(other)
+        return _generic.add(other, self)
 
     def __sub__(self, other):
         return _generic.subtract(self, other)

diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc
@@ -191,6 +191,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const AddNode* op) {
     // canonicalization rule
     // will try rewrite again after canonicalization.
     TVM_TRY_RECURSIVE_REWRITE(x + (c1 - y), (x - y) + c1);
+    TVM_TRY_RECURSIVE_REWRITE((c1 - y) + x, (x - y) + c1);
     TVM_TRY_RECURSIVE_REWRITE(x + c1 + y, (x + y) + c1);
     TVM_TRY_RECURSIVE_REWRITE(x + (c1 + y), (x + y) + c1);
     TVM_TRY_RECURSIVE_REWRITE(x + max(y, z), max(y, z) + x);

diff --git a/tests/python/unittest/test_arith_rewrite_simplify.py b/tests/python/unittest/test_arith_rewrite_simplify.py
@@ -184,6 +184,16 @@ def test_add_index_simplify():
     ck.verify(y * x + 10 * x, x * (y + 10))
     ck.verify(x * y + 10 * x, x * (y + 10))
 
+    ck.verify((2 * z) + tvm.te.min(x, y - (2 * z)), tvm.te.min(x + (z * 2), y))
+    ck.verify(y * x + x, x * (y + 1))
+    ck.verify(x * y + x, x * (y + 1))
+    ck.verify((x + 10) + 13, x + 23)
+    ck.verify((x + 10) + (13 + z), x + z + 23)
+    ck.verify(x * y + 10 * x, x * (y + 10))
+    ck.verify(y * x + x * 3, x * (y + 3))
+    ck.verify(x + 3 + y, x + y + 3)
+    ck.verify((3 - y) + x, x - y + 3)
+
 
     # canonicalization
     ck.verify(x + 2 + 3 + 4 + x, x * 2 + 9);