[BACKPORT-0.6][Bugfix][Arith] keep div_mode during floordiv simplify (a…

…pache#5922)
trevor-m · Jun 30, 2020 · 9712fef · 9712fef
1 parent 8d151ac
commit 9712fef
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/arith/canonical_simplify.cc b/src/arith/canonical_simplify.cc
@@ -692,6 +692,7 @@ SplitExpr CanonicalSimplifier::Impl::SplitDivConst(SplitExpr lhs, int64_t cval,
   CHECK(lhs->DivModeCompatibleTo(div_mode));
   CHECK_EQ(lhs->scale, 1);
   lhs.CopyOnWrite()->lower_factor *= cval;
+  lhs.CopyOnWrite()->div_mode = div_mode;
   return lhs;
 }
 

diff --git a/tests/python/unittest/test_arith_canonical_simplify.py b/tests/python/unittest/test_arith_canonical_simplify.py
@@ -85,6 +85,7 @@ def test_split_index_simplify():
     # floordiv
     fld = tvm.te.floordiv
     flm = tvm.te.floormod
+    ck.verify(fld(x*5, 2), fld(x*5, 2))
     ck.verify(fld(x, 3) * 3 + flm(x, 3), x)
     ck.verify(fld(x, 6) * 6 + flm(fld(x, 3), 2) * 3 + flm(x, 3), x)
     ck.verify(fld(fld(flm(x, 16), 2) * 2, 4), fld(flm(x, 16), 4))
@@ -124,8 +125,8 @@ def test_floormod_simplify():
     ck = CanonicalChecker()
     flm = tvm.te.floormod
     x, y = te.var("x"), te.var("y")
-    ck.verify(flm(flm((x*4) + y  - 466036, 24528) - 24512,  16),
-              flm((x*4) + y  + 12, 16))
+    ck.verify(flm(flm((x*4) + y - 466036, 24528) - 24512,  16),
+              flm((x*4) + y + 12, 16))
     ck.verify(flm(flm((x*4), 16), 8), flm(x, 2) * 4)