address review comments, fix div sub rule

apache · Jul 8, 2019 · 9cf0140 · 9cf0140
1 parent 36247ab
commit 9cf0140
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/src/arithmetic/int_operator.h b/src/arithmetic/int_operator.h
@@ -20,7 +20,7 @@
 /*!
  *  Copyright (c) 2019 by Contributors
  * \file int_operator.h
- * \brief Additional useful operators for intteger.
+ * \brief Additional useful operators for integer.
  */
 #ifndef TVM_ARITHMETIC_INT_OPERATOR_H_
 #define TVM_ARITHMETIC_INT_OPERATOR_H_

diff --git a/src/arithmetic/rewrite_simplify.cc b/src/arithmetic/rewrite_simplify.cc
@@ -395,10 +395,10 @@ Mutate_(const Sub* op, const Expr& self) {
                        c3.Eval()->value == c1.Eval()->value * c2.Eval()->value);
 
     TVM_TRY_REWRITE_IF(floordiv(x + c1, c3) - floordiv(x + c2, c3),
-                       floordiv(floormod(x + floormod(c1, c3), c3) + (c1 - c2), c3),
+                       floordiv(floormod(x + floormod(c2, c3), c3) + (c1 - c2), c3),
                        c3.Eval()->value > 0);
     TVM_TRY_REWRITE_IF(floordiv(x + c1, c3)  - floordiv(x, c3),
-                       floordiv(floormod(x + floormod(c1, c3), c3) + c1, c3),
+                       floordiv(floormod(x, c3) + c1, c3),
                        c3.Eval()->value > 0);
 
     // canonicalization rule

diff --git a/tests/python/unittest/test_arith_rewrite_simplify.py b/tests/python/unittest/test_arith_rewrite_simplify.py
@@ -288,7 +288,8 @@ def test_sub_index_simplify():
     ck.analyzer.update(x, tvm.arith.ConstIntBound(-1000, 1000), override=True)
     ck.analyzer.update(y, tvm.arith.ConstIntBound(-1000, 1000), override=True)
     ck.verify(x - fld(x, 3) * 3, flm(x, 3))
-    ck.verify(fld(x + 5, 3) - fld(x, 3), fld(flm(x + 2, 3) + 5, 3))
+    ck.verify(fld(x + 5, 3) - fld(x, 3), fld(flm(x, 3) + 5, 3))
+    ck.verify(fld(x + 5, 3) - fld(x + 2, 3), fld(flm(x + 2, 3), 3) + 1)
 
     ck.verify(fld(y, 3) * 3 - y, 0 - flm(y, 3))
     ck.verify(y - fld(y - 6, 5) * 5, flm(y + (-6), 5) + 6)