[Mlir-commits] [mlir] [MLIR][Affine] Simplify affine.for bounds by pruning redundant expressions via ValueBoundsConstraintSet (PR #199032)

[Oleksandr Alex Zinenko]

================
@@ -218,11 +218,105 @@ struct SimplifyDelinearizeOfLinearizeDisjoint final
   }
 };
 
+/// Simplifies the affine map results by eliminating redundant expressions.
+///
+/// This function performs a pairwise comparison of all expressions in the map
+/// using the analysis from `ValueBoundsConstraintSet`. If an expression `a` is
+/// statically proven to be strictly bounded or covered by another expression
+/// `b` (based on the given comparison operator `cmp`), `a` is considered
+/// redundant and is safely pruned from the results. NOTE: We iterate backwards
+/// (from size-1 down to 0) to safely erase elements from the `SmallVector`
+/// without causing iterator invalidation or indexing shifts for upcoming
+/// elements.
+static SmallVector<AffineExpr>
+simplifyRedundantMapResults(AffineMap map, SmallVector<Value> operands,
+                            ValueBoundsConstraintSet::ComparisonOperator cmp) {
+  SmallVector<AffineExpr> mapResults(map.getResults());
+  auto *context = map.getContext();
+
+  for (int i = mapResults.size() - 1; i >= 0; --i) {
+    AffineExpr a = mapResults[i];
+    AffineMap mapA =
+        AffineMap::get(map.getNumDims(), map.getNumSymbols(), a, context);
+    ValueBoundsConstraintSet::Variable varA(mapA, operands);
+    bool shouldErase = false;
+
+    for (int j = 0, e = mapResults.size(); j < e; ++j) {
+      if (i == j)
+        continue;
+
+      AffineExpr b = mapResults[j];
+      AffineMap mapB =
+          AffineMap::get(map.getNumDims(), map.getNumSymbols(), b, context);
+      ValueBoundsConstraintSet::Variable varB(mapB, operands);
+
+      if (ValueBoundsConstraintSet::compare(varB, cmp, varA)) {
+        shouldErase = true;
+        break;
+      }
+    }
+
+    if (shouldErase)
+      mapResults.erase(mapResults.begin() + i);
----------------
ftynse wrote:

A better approach is to iterate over the map (use getSubMap to produce single-expression maps) and collect indices at which to earse the expressions. Then create a new vector of expressions and copy those expressions whose index wasn't marked for erasure. This avoids additional copies of elements every time an element is earsed.

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