[Mlir-commits] [mlir] dd7d81e - Fix simplification of x + x//c*-c to x mod c. (#98909)

[llvmlistbot at llvm.org]

Author: Johannes Reifferscheid
Date: 2024-07-15T16:59:47+02:00
New Revision: dd7d81ea49bf39e1d69bbb84bd3f31bd95519369

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

LOG: Fix simplification of x + x//c*-c to x mod c. (#98909)

There was no check that rhs is actually a multiplication.

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 bfb7c4849356e..75cc01ee9a146 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -751,8 +751,10 @@ static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
   }
 
   // Process lrhs, which is 'expr floordiv c'.
+  // expr + (expr // c * -c) = expr % c
   AffineBinaryOpExpr lrBinOpExpr = dyn_cast<AffineBinaryOpExpr>(lrhs);
-  if (!lrBinOpExpr || lrBinOpExpr.getKind() != AffineExprKind::FloorDiv)
+  if (!lrBinOpExpr || rhs.getKind() != AffineExprKind::Mul ||
+      lrBinOpExpr.getKind() != AffineExprKind::FloorDiv)
     return nullptr;
 
   llrhs = lrBinOpExpr.getLHS();

diff  --git a/mlir/unittests/IR/AffineExprTest.cpp b/mlir/unittests/IR/AffineExprTest.cpp
index a0affc4341b0b..75c893334943d 100644
--- a/mlir/unittests/IR/AffineExprTest.cpp
+++ b/mlir/unittests/IR/AffineExprTest.cpp
@@ -98,3 +98,11 @@ TEST(AffineExprTest, divisionSimplification) {
   ASSERT_EQ((d0 * 6).ceilDiv(4).getKind(), AffineExprKind::CeilDiv);
   ASSERT_EQ((d0 * 6).ceilDiv(-2), d0 * -3);
 }
+
+TEST(AffineExprTest, modSimplificationRegression) {
+  MLIRContext ctx;
+  OpBuilder b(&ctx);
+  auto d0 = b.getAffineDimExpr(0);
+  auto sum = d0 + d0.floorDiv(3).floorDiv(-3);
+  ASSERT_EQ(sum.getKind(), AffineExprKind::Add);
+}