From 52a7c4a77526c4054b1003901c5de5e5a317562c Mon Sep 17 00:00:00 2001 From: Wang Yao Date: Sun, 3 May 2020 17:28:53 -0700 Subject: [PATCH] Fix Canonical Simplifier --- src/arith/canonical_simplify.cc | 1 + tests/python/unittest/test_arith_canonical_simplify.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/arith/canonical_simplify.cc b/src/arith/canonical_simplify.cc index 2bb0189890a6..a10db7a92020 100644 --- a/src/arith/canonical_simplify.cc +++ b/src/arith/canonical_simplify.cc @@ -870,6 +870,7 @@ SplitModConst(SplitExpr lhs, int64_t cval, DivMode div_mode) { lhs->upper_factor != SplitExprNode::kPosInf) { auto updated = ToSplitExpr(this->VisitExpr(ModImpl( lhs->index, make_const(lhs.dtype(), new_upper_factor), div_mode))); + updated.CopyOnWrite()->scale = lhs->scale; // re-apply the lower_factor if (lhs->lower_factor != 1) { return SplitDivConst(updated, lhs->lower_factor, div_mode); diff --git a/tests/python/unittest/test_arith_canonical_simplify.py b/tests/python/unittest/test_arith_canonical_simplify.py index 0dcf1fb5344c..179152273c00 100644 --- a/tests/python/unittest/test_arith_canonical_simplify.py +++ b/tests/python/unittest/test_arith_canonical_simplify.py @@ -126,7 +126,7 @@ def test_floormod_simplify(): 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), 16), 8), flm(x, 2) * 4) def test_canonical_mixed():