[Mlir-commits] [mlir] [MLIR][Affine] Keep dims when simplifying affine.min/max (PR #235819)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: SunXu (whutsunxu)

<details>
<summary>Changes</summary>

## Summary
- `simplifyAffineMinMaxOp` rebuilt maps as 0-dim / all-symbol, so a loop IV was used as a symbol and `affine.apply` verify failed (https://github.com/llvm/llvm-project/issues/235064).
- Call `canonicalizeMapAndOperands` after rebuilding so demoted dims are legalized.
- Add dim and symbol `affine.min` tests in `simplify-min-max-ops.mlir`.

## Test plan
- [x] `llvm-lit` `--filter='Dialect/Affine/simplify-min-max-ops.mlir'`
- [x] `mlir-opt` on the issue 235064 repro with `--verify-each`

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/SimplifyAffineMinMax.cpp (+6-2) 
- (modified) mlir/test/Dialect/Affine/simplify-min-max-ops.mlir (+23) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineMinMax.cpp b/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineMinMax.cpp
index ea7f3157c82eaa..88ffbd1461ce55 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineMinMax.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineMinMax.cpp
@@ -35,7 +35,7 @@ static bool simplifyAffineMinMaxOp(RewriterBase &rewriter, AffineOp affineOp) {
   using ComparisonOperator = ValueBoundsConstraintSet::ComparisonOperator;
 
   AffineMap affineMap = affineOp.getMap();
-  ValueRange operands = affineOp.getOperands();
+  SmallVector<Value> operands(affineOp.getOperands());
   static constexpr bool isMin = std::is_same_v<AffineOp, AffineMinOp>;
 
   LDBG() << "analyzing value: `" << affineOp;
@@ -131,9 +131,13 @@ static bool simplifyAffineMinMaxOp(RewriterBase &rewriter, AffineOp affineOp) {
   affineMap =
       AffineMap::get(0, affineMap.getNumSymbols() + affineMap.getNumDims(),
                      results, rewriter.getContext());
+  canonicalizeMapAndOperands(&affineMap, &operands);
 
   // Update the affine op.
-  rewriter.modifyOpInPlace(affineOp, [&]() { affineOp.setMap(affineMap); });
+  rewriter.modifyOpInPlace(affineOp, [&]() {
+    affineOp.setMap(affineMap);
+    affineOp->setOperands(operands);
+  });
   LDBG() << "- simplified affine op: `" << affineOp << "`";
   return true;
 }
diff --git a/mlir/test/Dialect/Affine/simplify-min-max-ops.mlir b/mlir/test/Dialect/Affine/simplify-min-max-ops.mlir
index 96e3c32f6d49fe..b17a66d5cecb22 100644
--- a/mlir/test/Dialect/Affine/simplify-min-max-ops.mlir
+++ b/mlir/test/Dialect/Affine/simplify-min-max-ops.mlir
@@ -3,6 +3,7 @@
 // CHECK-DAG: #[[MAP_0:.*]] = affine_map<()[s0] -> (32, s0)>
 // CHECK-DAG: #[[MAP_1:.*]] = affine_map<()[s0, s1] -> (s1, s0)>
 // CHECK-DAG: #[[MAP_2:.*]] = affine_map<()[s0] -> (256, s0)>
+// CHECK-DAG: #[[MAP_MOD:.*]] = affine_map<(d0) -> (-d0 + (d0 floordiv 2) * 2)>
 
 // CHECK: @min_max_full_simplify
 func.func @min_max_full_simplify() -> (index, index) {
@@ -58,3 +59,25 @@ func.func @overlapping_constraints() -> (index, index) {
   %r1 = affine.max affine_map<()[s0, s1, s2] -> (s0, s1, s2)>()[%0, %1, %2]
   return %r0, %r1 : index, index
 }
+
+// CHECK: @affine_min_dim_simplify
+func.func @affine_min_dim_simplify(%bound: index) -> index {
+  // CHECK: affine.for %[[IV:.*]] = 0 to
+  %out = affine.for %iv = 0 to %bound iter_args(%carried = %bound) -> (index) {
+    // CHECK: affine.apply #[[MAP_MOD]](%[[IV]])
+    %result = affine.min affine_map<(d0) -> (0, -(d0 mod 2))>(%iv)
+    affine.yield %result : index
+  }
+  return %out : index
+}
+
+// CHECK: @affine_min_symbol_simplify
+func.func @affine_min_symbol_simplify(%bound: index) -> index {
+  // CHECK: affine.for %[[IV:.*]] = 0 to
+  %out = affine.for %iv = 0 to %bound iter_args(%carried = %bound) -> (index) {
+    // CHECK: affine.apply #[[MAP_MOD]](%[[IV]])
+    %result = affine.min affine_map<()[s0] -> (0, -(s0 mod 2))>()[%iv]
+    affine.yield %result : index
+  }
+  return %out : index
+}

``````````

</details>


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