[Mlir-commits] [mlir] 8272b6b - [mlir][IR] Fix bug in AffineExpr simplifier `lhs % rhs` where `lhs = lhs floordiv rhs` (#119245)

[llvmlistbot at llvm.org]

Author: Christopher Bate
Date: 2024-12-18T15:22:57-07:00
New Revision: 8272b6bd6146aab973ff7018ad642b99fde00904

URL: https://github.com/llvm/llvm-project/commit/8272b6bd6146aab973ff7018ad642b99fde00904
DIFF: https://github.com/llvm/llvm-project/commit/8272b6bd6146aab973ff7018ad642b99fde00904.diff

LOG: [mlir][IR] Fix bug in AffineExpr simplifier `lhs % rhs` where `lhs = lhs floordiv rhs` (#119245)

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 https://github.com/llvm/llvm-project/issues/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.

Resolves https://github.com/llvm/llvm-project/issues/114654.

Added: 
    

Modified: 
    mlir/lib/IR/AffineExpr.cpp
    mlir/unittests/IR/AffineExprTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 2291d64c50a560..59df0cd6833db3 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/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();
 }

diff  --git a/mlir/unittests/IR/AffineExprTest.cpp b/mlir/unittests/IR/AffineExprTest.cpp
index 9e89a5b79e2e2e..8a2d697540d5ca 100644
--- a/mlir/unittests/IR/AffineExprTest.cpp
+++ b/mlir/unittests/IR/AffineExprTest.cpp
@@ -129,3 +129,21 @@ 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);
+
+  // Manually replace variables by constants to avoid constant folding.
+  AffineExpr expr = (d0 - (d0 + 2)).floorDiv(8) % 8;
+  AffineExpr result = mlir::simplifyAffineExpr(expr, 1, 0);
+
+  ASSERT_TRUE(isa<AffineConstantExpr>(result));
+  ASSERT_EQ(cast<AffineConstantExpr>(result).getValue(), 7);
+}