[Mlir-commits] [mlir] [MLIR] Fix bug in AffineExpr simplifyAdd (PR #181613)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Uday Bondhugula (bondhugula)

<details>
<summary>Changes</summary>

Fix bug in AffineExpr simplifyAdd. This was leading to an incorrect construction of affine expressions in certain cases; the expression was being wrongly simplified to a modulo one.

Test cases for affine expr simplification by construction in lib/IR/AffineExpr.cpp were earlier added to a transforms pass' tests; add basic ones to test/IR/ while on this.

---
Full diff: https://github.com/llvm/llvm-project/pull/181613.diff


2 Files Affected:

- (modified) mlir/lib/IR/AffineExpr.cpp (+4-4) 
- (modified) mlir/test/IR/affine-map.mlir (+21) 


``````````diff
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index b476309e6bdfd..1cb3163c4ba94 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -738,7 +738,7 @@ static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
   // and readable form.
 
   // Process '(expr floordiv c) * (-c)'.
-  if (!rBinOpExpr)
+  if (!rBinOpExpr || rBinOpExpr.getKind() != AffineExprKind::Mul)
     return nullptr;
 
   auto lrhs = rBinOpExpr.getLHS();
@@ -746,10 +746,10 @@ static AffineExpr simplifyAdd(AffineExpr lhs, AffineExpr rhs) {
 
   AffineExpr llrhs, rlrhs;
 
-  // Check if lrhsBinOpExpr is of the form (expr floordiv q) * q, where q is a
-  // symbolic expression.
+  // Check if lrhsBinOpExpr is of the form (expr floordiv q) * q,
+  // where q is a symbolic expression.
   auto lrhsBinOpExpr = dyn_cast<AffineBinaryOpExpr>(lrhs);
-  // Check rrhsConstOpExpr = -1.
+  // Check rrhsConstOpExpr = -1 as part of ((expr floordiv q) * q)) * (-1).
   auto rrhsConstOpExpr = dyn_cast<AffineConstantExpr>(rrhs);
   if (rrhsConstOpExpr && rrhsConstOpExpr.getValue() == -1 && lrhsBinOpExpr &&
       lrhsBinOpExpr.getKind() == AffineExprKind::Mul) {
diff --git a/mlir/test/IR/affine-map.mlir b/mlir/test/IR/affine-map.mlir
index 6277b28561f36..129e83535b4a4 100644
--- a/mlir/test/IR/affine-map.mlir
+++ b/mlir/test/IR/affine-map.mlir
@@ -207,6 +207,18 @@
 // CHECK: #map{{[0-9]*}} = affine_map<(d0, d1)[s0] -> (d0 + d1 + s0)>
 #map64 = affine_map<(i0, i1)[mod] -> (i0 + i1 + mod)>
 
+// Simplifying add expressions with floordiv on RHS to a mod.
+// CHECK: #map{{[0-9]*}} = affine_map<(d0) -> (d0 mod 4)>
+#map65 = affine_map<(d0) -> (d0 - d0 floordiv 4 * 4)>
+
+// Same as above but with semi-affine.
+// CHECK: #map{{[0-9]*}} = affine_map<(d0)[s0] -> (d0 mod (s0 * 4))>
+#map65a = affine_map<(d0)[s0] -> (d0 - d0 floordiv (4 * s0) * (4 * s0))>
+
+// Negative test for the above; shouldn't get simplified.
+// CHECK: #map{{[0-9]*}} = affine_map<(d0) -> (d0 + d0 floordiv 4 - 4)>
+#map66 = affine_map<(d0) -> (d0 + ((d0 floordiv 4) - 4))>
+
 // Single identity maps are removed.
 // CHECK: @f0(memref<2x4xi8, 1>)
 func.func private @f0(memref<2x4xi8, #map0, 1>)
@@ -409,3 +421,12 @@ func.func private @f56(memref<1x1xi8, #map56>)
 
 // CHECK: "f64"() {map = #map{{[0-9]*}}} : () -> ()
 "f64"() {map = #map64} : () -> ()
+
+// CHECK: "f65"() {map = #map{{[0-9]*}}} : () -> ()
+"f65"() {map = #map65} : () -> ()
+
+// CHECK: "f65a"() {map = #map{{[0-9]*}}} : () -> ()
+"f65a"() {map = #map65a} : () -> ()
+
+// CHECK: "f66"() {map = #map{{[0-9]*}}} : () -> ()
+"f66"() {map = #map66} : () -> ()

``````````

</details>


https://github.com/llvm/llvm-project/pull/181613