[Mlir-commits] [mlir] 8af0860 - AffineExpr: Fix result of d0 + (d0 // -c) * c. (#107530)

[llvmlistbot at llvm.org]

Author: Johannes Reifferscheid
Date: 2024-09-06T12:53:33+02:00
New Revision: 8af0860529d75b61b66cb96ac05f503b0e2e845c

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

LOG: AffineExpr: Fix result of d0 + (d0 // -c) * c. (#107530)

Currently, this is rewritten to d0 mod -c. However, we do not support
modulo with a negative RHS in our lowering passes, so this triggers
undefined behavior.

It would be better to not have these ad hoc simplifications at all, but
I guess that ship has sailed.

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 fc7ede279643ed..0b078966aeb85b 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -760,8 +760,11 @@ static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
 
   llrhs = lrBinOpExpr.getLHS();
   rlrhs = lrBinOpExpr.getRHS();
+  auto rlrhsConstOpExpr = dyn_cast<AffineConstantExpr>(rlrhs);
+  // We don't support modulo with a negative RHS.
+  bool isPositiveRhs = rlrhsConstOpExpr && rlrhsConstOpExpr.getValue() > 0;
 
-  if (lhs == llrhs && rlrhs == -rrhs) {
+  if (isPositiveRhs && lhs == llrhs && rlrhs == -rrhs) {
     return lhs % rlrhs;
   }
   return nullptr;

diff  --git a/mlir/unittests/IR/AffineExprTest.cpp b/mlir/unittests/IR/AffineExprTest.cpp
index dc78bbac85f3cf..9e89a5b79e2e2e 100644
--- a/mlir/unittests/IR/AffineExprTest.cpp
+++ b/mlir/unittests/IR/AffineExprTest.cpp
@@ -15,6 +15,13 @@
 
 using namespace mlir;
 
+static std::string toString(AffineExpr expr) {
+  std::string s;
+  llvm::raw_string_ostream ss(s);
+  ss << expr;
+  return s;
+}
+
 // Test creating AffineExprs using the overloaded binary operators.
 TEST(AffineExprTest, constructFromBinaryOperators) {
   MLIRContext ctx;
@@ -112,3 +119,13 @@ TEST(AffineExprTest, divisorOfNegativeFloorDiv) {
   OpBuilder b(&ctx);
   ASSERT_EQ(b.getAffineDimExpr(0).floorDiv(-1).getLargestKnownDivisor(), 1);
 }
+
+TEST(AffineExprTest, d0PlusD0FloorDivNeg2) {
+  // Regression test for a bug where this was rewritten to d0 mod -2. We do not
+  // support a negative RHS for mod in LowerAffinePass.
+  MLIRContext ctx;
+  OpBuilder b(&ctx);
+  auto d0 = b.getAffineDimExpr(0);
+  auto sum = d0 + d0.floorDiv(-2) * 2;
+  ASSERT_EQ(toString(sum), "d0 + (d0 floordiv -2) * 2");
+}