diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp index bfb7c4849356eb..75cc01ee9a146a 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(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 a0affc4341b0b4..75c893334943d3 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); +}