[Mlir-commits] [mlir] 3986949 - [MLIR][Affine] Simplify affine.for bounds by pruning redundant expressions via ValueBoundsConstraintSet (#199032)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 10 07:00:21 PDT 2026
Author: lonely eagle
Date: 2026-06-10T22:00:15+08:00
New Revision: 398694988742d26a5eaca038c337c193d25c0268
URL: https://github.com/llvm/llvm-project/commit/398694988742d26a5eaca038c337c193d25c0268
DIFF: https://github.com/llvm/llvm-project/commit/398694988742d26a5eaca038c337c193d25c0268.diff
LOG: [MLIR][Affine] Simplify affine.for bounds by pruning redundant expressions via ValueBoundsConstraintSet (#199032)
This PR introduces a new pattern `SimplifyAffineForBoundMap` to simplify
multi-result lower and upper bounds of `affine.for` loops by pruning
redundant expressions leveraging `ValueBoundsConstraintSet`. This PR
addresses the limitation where the built-in folder of `affine.for` can
only handle constant bounds.
Added:
Modified:
mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp
mlir/test/Dialect/Affine/simplify-with-bounds.mlir
Removed:
################################################################################
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..f0c16a05a3070 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SimplifyAffineWithBounds.cpp
@@ -218,11 +218,96 @@ 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.
+static SmallVector<AffineExpr>
+simplifyRedundantMapResults(AffineMap map, ValueRange operands,
+ ValueBoundsConstraintSet::ComparisonOperator cmp) {
+ llvm::BitVector preservedExprs(map.getNumResults(), true);
+ for (size_t i = 0, e = map.getNumResults(); i < e; ++i) {
+ AffineMap mapA = map.getSubMap(i);
+ ValueBoundsConstraintSet::Variable varA(mapA, operands);
+
+ for (size_t j = 0; j < e; ++j) {
+ if (i == j || !preservedExprs[j])
+ continue;
+
+ AffineMap mapB = map.getSubMap(j);
+ ValueBoundsConstraintSet::Variable varB(mapB, operands);
+
+ if (ValueBoundsConstraintSet::compare(varB, cmp, varA)) {
+ preservedExprs[i] = false;
+ break;
+ }
+ }
+ }
+
+ SmallVector<AffineExpr> mapResults;
+ for (size_t i = 0, e = map.getNumResults(); i < e; ++i)
+ if (preservedExprs[i])
+ mapResults.push_back(map.getResult(i));
+ return mapResults;
+}
+
+/// A 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();
+
+ MLIRContext *context = forOp->getContext();
+ if (lowerBoundUpdate) {
+ rewriter.modifyOpInPlace(forOp, [&]() {
+ forOp.setLowerBound(forOp.getLowerBoundOperands(),
+ AffineMap::get(lowerBoundMap.getNumDims(),
+ lowerBoundMap.getNumSymbols(),
+ lowerMapExprs, context));
+ });
+ }
+ if (upperBoundUpdate) {
+ rewriter.modifyOpInPlace(forOp, [&]() {
+ forOp.setUpperBound(forOp.getUpperBoundOperands(),
+ AffineMap::get(upperBoundMap.getNumDims(),
+ upperBoundMap.getNumSymbols(),
+ upperMapExprs, context));
+ });
+ }
+ 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: }
More information about the Mlir-commits
mailing list