This PR fixes a bug revealed as part of the IterMapSimplify change.

y - (x % 2) were simplified to y + (x % 2) - 1 which is wrong.

Regression tests are added.

Author: tqchen

src/arith/iter_affine_map.cc

@@ -1093,11 +1093,6 @@ class IterMapRewriter : public ExprMutator {
   PrimExpr SplitFloorModConst(IterSplitExpr lhs, PrimExpr base, PrimExpr rhs);
 
   static void AddToLhs(IterSumExprNode* lhs, IterSplitExpr rhs, int sign) {
-    if (sign < 0 && is_const_int(rhs->extent, 2)) {
-      lhs->base -= rhs->scale;
-      sign = 1;
-    }
-
     tir::ExprDeepEqual equal;
     for (size_t i = 0; i < lhs->args.size(); ++i) {
       IterSplitExpr lvalue = lhs->args[i];

tests/python/unittest/test_arith_iter_affine_map.py

@@ -222,18 +222,6 @@ def test_compound():
     assert_iter_sum_pattern({z[0]: (18, 0, 1, sz), xi[0]: (5, 0)}, var_dom([(x, 10), (y, 9)]))
 
 
-def test_compound_floormod_two():
-    x = tvm.tir.Var("x", "int32")
-    fld = tvm.tir.floordiv
-    flm = tvm.tir.floormod
-
-    # extent of 2 are normalized to positive scale
-    assert_iter_sum_pattern(
-        expect_dict={fld(x, 2) * 2 - flm(x, 2) + 1: (8, 0, 1)},
-        dom_map=var_dom([(x, 8)]),
-    )
-
-
 def test_predicate():
     x = tvm.tir.Var("x", "int32")
     y = tvm.tir.Var("y", "int32")
@@ -1190,6 +1178,12 @@ def test_iter_map_simplify_unit_loop_order():
         simplify_trivial_iterators=False,
     )
 
+    assert_iter_map_simplfy(
+        {y + 64 - x % 2 * 64: y + 64 - x % 2 * 64},
+        var_dom([(x, 6), (y, 64)]),
+        simplify_trivial_iterators=False,
+    )
+
 
 if __name__ == "__main__":
     tvm.testing.main()