[Mlir-commits] [mlir] [MLIR][Affine] Simplify affine.for bounds by pruning redundant expressions via ValueBoundsConstraintSet (PR #199032)
lonely eagle
llvmlistbot at llvm.org
Thu May 21 23:10:52 PDT 2026
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/199032
>From f994616e7c2bd969e7ea6aa525e286d1b5476695 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Thu, 21 May 2026 13:53:15 +0000
Subject: [PATCH 1/2] add SimplifyAffineLoopBoundMap pattern.
---
.../mlir/Dialect/Affine/Transforms/Passes.td | 4 +-
.../Transforms/SimplifyAffineWithBounds.cpp | 97 ++++++++++++++++++-
.../Dialect/Affine/simplify-with-bounds.mlir | 30 +++++-
3 files changed, 127 insertions(+), 4 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
index 1c54ca7deca91..7cecbcac5b932 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
@@ -434,8 +434,8 @@ def SimplifyAffineWithBounds : Pass<"affine-simplify-with-bounds"> {
let summary = "Simplify affine index operations using value bounds analysis";
let description = [{
This pass simplifies `affine.delinearize_index` / `affine.linearize_index`
- pairs by using value bounds analysis to match basis products. Unlike the
- built-in canonicalization patterns which only use exact `OpFoldResult`
+ pairs and `affine.for` by using value bounds analysis to match basis products.
+ Unlike the built-in canonicalization patterns which only use exact `OpFoldResult`
comparisons, this pass can prove equality of dynamic basis products through
`ValueBoundsConstraintSet`.
}];
diff --git a/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp
index 6e7d5d91334c0..00b4a2b3c1c55 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp
@@ -218,11 +218,106 @@ 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);
+ }
+
+ return mapResults;
+}
+
+/// A canonicalization pattern that simplifies multi-result lower and upper
+/// bounds of `affine.for` loops by pruning redundant expressions leveraging
+/// `ValueBoundsConstraintSet`.
+struct SimplifyAffineLoopBoundMap final : OpRewritePattern<AffineForOp> {
+ using Base::Base;
+ LogicalResult matchAndRewrite(AffineForOp forOp,
+ PatternRewriter &rewriter) const override {
+ AffineMap lowerBoundMap = forOp.getLowerBoundMap();
+ auto lowerBoundOperands = forOp.getLowerBoundOperands();
+ AffineMap upperBoundMap = forOp.getUpperBoundMap();
+ auto upperBoundOperands = forOp.getUpperBoundOperands();
+ if (lowerBoundMap.getNumResults() < 2 &&
+ forOp.getUpperBoundMap().getNumResults() < 2)
+ return failure();
+
+ SmallVector<AffineExpr> lowerMapExprs = simplifyRedundantMapResults(
+ lowerBoundMap, lowerBoundOperands, ValueBoundsConstraintSet::GT);
+ SmallVector<AffineExpr> upperMapExprs = simplifyRedundantMapResults(
+ upperBoundMap, upperBoundOperands, ValueBoundsConstraintSet::LT);
+
+ bool lowerBoundUpdate =
+ lowerMapExprs.size() < lowerBoundMap.getNumResults();
+ bool upperBoundUpdate =
+ upperMapExprs.size() < upperBoundMap.getNumResults();
+ if (!(lowerBoundUpdate || upperBoundUpdate))
+ return failure();
+
+ auto *context = forOp->getContext();
+ if (lowerBoundUpdate) {
+ AffineMap newMap =
+ AffineMap::get(lowerBoundMap.getNumDims(),
+ lowerBoundMap.getNumSymbols(), lowerMapExprs, context);
+ rewriter.modifyOpInPlace(forOp, [&]() {
+ forOp.setLowerBound(forOp.getLowerBoundOperands(), newMap);
+ });
+ }
+ if (upperBoundUpdate) {
+ AffineMap newMap =
+ AffineMap::get(upperBoundMap.getNumDims(),
+ upperBoundMap.getNumSymbols(), upperMapExprs, context);
+ rewriter.modifyOpInPlace(forOp, [&]() {
+ forOp.setUpperBound(forOp.getUpperBoundOperands(), newMap);
+ });
+ }
+ return success();
+ }
+};
} // namespace
void affine::populateSimplifyAffineWithBoundsPatterns(
RewritePatternSet &patterns) {
- patterns.add<SimplifyDelinearizeOfLinearizeDisjoint>(patterns.getContext());
+ patterns
+ .add<SimplifyDelinearizeOfLinearizeDisjoint, SimplifyAffineLoopBoundMap>(
+ patterns.getContext());
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Affine/simplify-with-bounds.mlir b/mlir/test/Dialect/Affine/simplify-with-bounds.mlir
index a9c517a5c5add..6cf474076f6fa 100644
--- a/mlir/test/Dialect/Affine/simplify-with-bounds.mlir
+++ b/mlir/test/Dialect/Affine/simplify-with-bounds.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -affine-simplify-with-bounds %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -affine-simplify-with-bounds %s | FileCheck %s
// CHECK-LABEL: func @many_to_one_static_tail
// CHECK-SAME: %[[A:.*]]: index, %[[B:.*]]: index, %[[C:.*]]: index
@@ -165,3 +165,31 @@ func.func @input_not_linearize(%x: index) -> (index, index) {
%0:2 = affine.delinearize_index %x into (4, 8) : index, index
return %0#0, %0#1 : index, index
}
+
+// -----
+
+// CHECK-DAG: #[[$MAP:.+]] = affine_map<()[s0, s1] -> (s0, s1)>
+
+// CHECK-LABEL: func @simplify_loop_bound
+func.func @simplify_loop_bound() -> index{
+ %c0 = arith.constant 0 :index
+ %c1 = arith.constant 1 : index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ %bound1 = test.value_with_bounds { min = 2 : index, max = 3 : index}
+ %bound2 = test.value_with_bounds { min = 2 : index, max = 3 : index}
+ %res = affine.for %iv = max affine_map<(d0, d1, d2) -> (d0, d1,d2)>(%bound, %bound1,%bound2) to min affine_map<(d0, d1, d2) -> (d0, d1,d2)>(%bound, %bound1,%bound2) step 2 iter_args(%arg = %c0) -> index {
+ %sum = arith.addi %arg, %c1 : index
+ affine.yield %sum : index
+ }
+ return %res : index
+}
+
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[BOUND0:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
+// CHECK-DAG: %[[BOUND1:.*]] = test.value_with_bounds {max = 3 : index, min = 2 : index}
+// CHECK-DAG: %[[BOUND2:.*]] = test.value_with_bounds {max = 3 : index, min = 2 : index}
+// CHECK: affine.for %[[IV:.*]] = max #[[$MAP]]()[%[[BOUND1]], %[[BOUND2]]] to %[[BOUND0]] step 2 iter_args(%[[ARG:.*]] = %[[C0]]) -> (index) {
+// CHECK: %[[SUM:.*]] = arith.addi %[[ARG]], %[[C1]] : index
+// CHECK: affine.yield %[[SUM]] : index
+// CHECK: }
>From 9483a9385ea4fd220060ca41a8ea9509eb6d144d Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Fri, 22 May 2026 06:10:33 +0000
Subject: [PATCH 2/2] remove empty line.
---
mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp
index 00b4a2b3c1c55..a96432aebb9a5 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp
@@ -236,7 +236,6 @@ simplifyRedundantMapResults(AffineMap map, SmallVector<Value> operands,
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);
More information about the Mlir-commits
mailing list