[mlir][IR] Fix bug in AffineExpr simplifier lhs % rhs where lhs = lhs floordiv rhs

Fixes an issue where the `SimpleAffineExprFlattener` would simplify
`lhs % rhs` to just `-(lhs floordiv rhs)` instead of `lhs - (lhs floordiv rhs)`
if `lhs` happened to be equal to `lhs floordiv rhs`.

The reported failure case was `(d0, d1) -> (((d1 - (d1 + 2)) floordiv 8) % 8)`
from #114654.

Note that many paths that simplify AffineMaps (e.g. the AffineApplyOp
folder and canonicalization) would not observe this bug because of
of slightly different paths taken by the code. Slightly different grouping
of the terms could also result in avoiding the bug. The way to reproduce was
by constructing the map directly, replacing `d1` with `1` and calling
`mlir::simplifyAffineExpr`.

Resolves #114654.

Author: christopherbate

mlir/lib/IR/AffineExpr.cpp

@@ -1385,7 +1385,7 @@ LogicalResult SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) {
     lhs[getLocalVarStartIndex() + numLocals - 1] = -rhsConst;
   } else {
     // Reuse the existing local id.
-    lhs[getLocalVarStartIndex() + loc] = -rhsConst;
+    lhs[getLocalVarStartIndex() + loc] -= rhsConst;
   }
   return success();
 }

mlir/unittests/IR/AffineExprTest.cpp

@@ -129,3 +129,24 @@ TEST(AffineExprTest, d0PlusD0FloorDivNeg2) {
   auto sum = d0 + d0.floorDiv(-2) * 2;
   ASSERT_EQ(toString(sum), "d0 + (d0 floordiv -2) * 2");
 }
+
+TEST(AffineExprTEst, simpleAffineExprFlattenerRegression) {
+
+  // Regression test for a bug where mod simplification was not handled
+  // properly when `lhs % rhs` was happened to have the property that `lhs
+  // floordiv rhs = lhs`.
+  MLIRContext ctx;
+  OpBuilder b(&ctx);
+
+  auto d0 = b.getAffineDimExpr(0);
+  auto d1 = b.getAffineDimExpr(1);
+
+  // Manually replace variables by constants to avoid constant folding.
+  AffineExpr expr = (d0 - (d1 + 2)).floorDiv(8) % 8;
+  expr = expr.replaceDims(
+      {b.getAffineConstantExpr(1), b.getAffineConstantExpr(1)});
+  AffineExpr result = mlir::simplifyAffineExpr(expr, 2, 0);
+
+  ASSERT_TRUE(isa<AffineConstantExpr>(result));
+  ASSERT_EQ(cast<AffineConstantExpr>(result).getValue(), 7);
+}