[Mlir-commits] [mlir] b3115cd - [mlir][affine] Add useExpensiveMath option to AffineLoopNormalize pass (#211989)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 29 06:13:33 PDT 2026
Author: lonely eagle
Date: 2026-07-29T21:13:28+08:00
New Revision: b3115cd807196e465ce3fc16f0cab58be24e47fc
URL: https://github.com/llvm/llvm-project/commit/b3115cd807196e465ce3fc16f0cab58be24e47fc
DIFF: https://github.com/llvm/llvm-project/commit/b3115cd807196e465ce3fc16f0cab58be24e47fc.diff
LOG: [mlir][affine] Add useExpensiveMath option to AffineLoopNormalize pass (#211989)
This PR integrates constant bound inference into the existing
`AffineLoopNormalize` pass under a new `useExpensiveMath` option. When
`useExpensiveMath` is enabled (disabled by default due to potential
compilation time overhead), the pass leverages
`ValueBoundsConstraintSet` (uss presburger) analysis to refine dynamic
`affine.for` loop bounds into compile-time constant bounds. RFC:
https://discourse.llvm.org/t/rfc-mlir-enable-dynamic-and-tighter-affine-unrolling-via-valueboundsconstraintset/91055/2
---------
Co-authored-by: Oleksandr "Alex" Zinenko <azinenko at amd.com>
Added:
Modified:
mlir/include/mlir/Dialect/Affine/Transforms/Passes.h
mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
mlir/test/Dialect/Affine/affine-loop-normalize.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.h b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.h
index 87a7e74fbd196..ce764eb750583 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.h
@@ -60,7 +60,8 @@ std::unique_ptr<OperationPass<func::FuncOp>> createRaiseMemrefToAffine();
/// `promoteSingleIter` is true, single iteration loops are promoted (i.e., the
/// loop is replaced by its loop body).
std::unique_ptr<OperationPass<func::FuncOp>>
-createAffineLoopNormalizePass(bool promoteSingleIter = false);
+createAffineLoopNormalizePass(bool promoteSingleIter = false,
+ bool useExpensiveMath = false);
/// Performs packing (or explicit copying) of accessed memref regions into
/// buffers in the specified faster memory space through either pointwise copies
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
index 7cecbcac5b932..f2008eeae9fcd 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
@@ -385,6 +385,8 @@ def AffineLoopNormalize : Pass<"affine-loop-normalize", "func::FuncOp"> {
let options = [
Option<"promoteSingleIter", "promote-single-iter", "bool",
/*default=*/"true", "Promote single iteration loops">,
+ Option<"useExpensiveMath", "use-expensive-math", "bool", /*default=*/"false",
+ "Use Presburger arithmetic to constantize loop upper bounds (higher analysis cost)">
];
}
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
index c74d147016237..0d464631ad439 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
@@ -10,11 +10,14 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/Dialect/Affine/Transforms/Passes.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "llvm/ADT/SmallVector.h"
namespace mlir {
namespace affine {
@@ -28,13 +31,80 @@ using namespace mlir::affine;
namespace {
+/// Computes the constant upper or lower bound for a given affine map expression
+/// and its operands, constrained by the specified type.
+static FailureOr<int64_t> computeConstantBound(AffineMap map,
+ ValueRange operands,
+ presburger::BoundType type) {
+ ValueBoundsConstraintSet::Variable var(map, operands);
+ return ValueBoundsConstraintSet::computeConstantBound(
+ type, var, nullptr, {/*closedUb=*/true, /*allowIntegerType=*/true});
+}
+
+/// Attempts to infer a static constant upper bound for the given normalized
+/// `affine.for` loop using Value Bounds Analysis. If the dynamic upper bound's
+/// range [upperMin, upperMax] is proven to be a single constant value (upperMin
+/// == upperMax), the upper bound is directly replaced with this constant.
+/// Otherwise, if upperMin > 0, the loop is split (peeled) into a static main
+/// loop with a constant upper bound (`upperMin`) and a residual tail loop
+/// iterating from `upperMin` to the original dynamic bound.
+static LogicalResult
+inferAffineLoopUpperConstantBound(RewriterBase &b, AffineForOp forOp,
+ bool promoteSingleIter = true) {
+ // The loop is normalized so we can expect its lower bound to be 0 and step to
+ // be 1
+ if (!forOp.hasConstantLowerBound() || forOp.getConstantLowerBound() != 0)
+ return failure();
+ if (forOp.getStepAsInt() != 1)
+ return failure();
+ if (forOp.getUpperBoundMap().getNumResults() > 1)
+ return failure();
+
+ // Infer the range [upperMin, upperMax] for the upper bound. We require a
+ // strictly positive minimum bound (upperMin > 0) to guarantee a safe,
+ // non-empty static trip count for the main loop.
+ FailureOr<int64_t> upperMin = computeConstantBound(
+ forOp.getUpperBoundMap(), forOp.getUpperBoundOperands(),
+ presburger::BoundType::LB);
+ FailureOr<int64_t> upperMax = computeConstantBound(
+ forOp.getUpperBoundMap(), forOp.getUpperBoundOperands(),
+ presburger::BoundType::UB);
+ if (failed(upperMin) || *upperMin <= 0)
+ return failure();
+
+ // The upper bound is dynamic within [upperMin, upperMax]. Split the loop into
+ // a static main loop (0 to upperMin) and a residual tail loop (upperMin to
+ // dynamic bound).
+ if (failed(upperMax) || *upperMax > *upperMin) {
+ b.setInsertionPoint(forOp);
+ AffineForOp clonedForOp = cast<AffineForOp>(b.clone(*forOp));
+ clonedForOp.setConstantUpperBound(*upperMin);
+ forOp.setConstantLowerBound(*upperMin);
+ forOp.getInitsMutable().assign(clonedForOp->getResults());
+ if (promoteSingleIter)
+ (void)promoteIfSingleIteration(clonedForOp);
+
+ return success();
+ }
+
+ // If upperMin == upperMax. The upper bound is proven to be a strict constant
+ // at compile time. Directly constantize the bound without peeling a tail
+ // loop.
+ forOp.setConstantUpperBound(*upperMin);
+ if (promoteSingleIter)
+ (void)promoteIfSingleIteration(forOp);
+ return success();
+}
+
/// Normalize affine.parallel ops so that lower bounds are 0 and steps are 1.
/// As currently implemented, this pass cannot fail, but it might skip over ops
/// that are already in a normalized form.
struct AffineLoopNormalizePass
: public affine::impl::AffineLoopNormalizeBase<AffineLoopNormalizePass> {
- explicit AffineLoopNormalizePass(bool promoteSingleIter) {
+ explicit AffineLoopNormalizePass(bool promoteSingleIter,
+ bool useExpensiveMath) {
this->promoteSingleIter = promoteSingleIter;
+ this->useExpensiveMath = useExpensiveMath;
}
void runOnOperation() override {
@@ -44,12 +114,31 @@ struct AffineLoopNormalizePass
else if (auto affineFor = dyn_cast<AffineForOp>(op))
(void)normalizeAffineFor(affineFor, promoteSingleIter);
});
+
+ // Infer and rewrite the upper bound into a compile-time constant for each
+ // loop.
+ if (useExpensiveMath) {
+ IRRewriter b(&getContext());
+ SmallVector<AffineForOp> loops;
+
+ // Collect target loops because `inferAffineLoopUpperConstantBound` may
+ // create new loops during processing.
+ // TODO: When running `normalizeAffineFor` with `promoteSingleIter=true`,
+ // there is currently no clean way to know if the loop was promoted. We
+ // can improve this in the future to avoid calling `walk` to pre-collect
+ // loops.
+ getOperation()->walk([&](AffineForOp forOp) { loops.push_back(forOp); });
+ for (AffineForOp forOp : loops)
+ (void)inferAffineLoopUpperConstantBound(b, forOp, promoteSingleIter);
+ }
}
};
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>>
-mlir::affine::createAffineLoopNormalizePass(bool promoteSingleIter) {
- return std::make_unique<AffineLoopNormalizePass>(promoteSingleIter);
+mlir::affine::createAffineLoopNormalizePass(bool promoteSingleIter,
+ bool useExpensiveMath) {
+ return std::make_unique<AffineLoopNormalizePass>(promoteSingleIter,
+ useExpensiveMath);
}
diff --git a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
index 7d90efec0c21b..879ccd798eba3 100644
--- a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
@@ -1,5 +1,7 @@
// RUN: mlir-opt %s -affine-loop-normalize -split-input-file | FileCheck %s
// RUN: mlir-opt %s -affine-loop-normalize='promote-single-iter=1' -split-input-file | FileCheck %s --check-prefix=PROMOTE-SINGLE-ITER
+// RUN: mlir-opt %s -affine-loop-normalize='use-expensive-math=1' -split-input-file | FileCheck %s --check-prefix=USE-EXPENSIVE-MATH
+// RUN: mlir-opt %s -affine-loop-normalize='use-expensive-math=1 promote-single-iter=1' -split-input-file | FileCheck %s --check-prefix=USE-EXPENSIVE-MATH-AND-PROMOTE
// Normalize steps to 1 and lower bounds to 0.
@@ -323,3 +325,97 @@ func.func @multi_level_tiled_matmul() {
}
return
}
+
+// -----
+
+// USE-EXPENSIVE-MATH-DAG: #[[$MAP_APPLY:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
+// USE-EXPENSIVE-MATH-DAG: #[[$MAP_UB:.+]] = affine_map<()[s0] -> ((-s0 + 9) ceildiv 2)>
+// USE-EXPENSIVE-MATH-LABEL: func @peeling_main_loop
+func.func @peeling_main_loop() {
+ %c0 = arith.constant 0 : index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ affine.for %iv = %bound to 9 step 2 iter_args(%arg = %c0) -> index {
+ %sum = arith.addi %arg, %bound : index
+ affine.yield %sum : index
+ }
+ return
+}
+
+// USE-EXPENSIVE-MATH: %[[C0:.*]] = arith.constant 0 : index
+// USE-EXPENSIVE-MATH: %[[BOUND:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
+// USE-EXPENSIVE-MATH: %[[MAIN_RES:.*]] = affine.for %[[IV_MAIN:.*]] = 0 to 4 iter_args(%[[ARG_MAIN:.*]] = %[[C0]]) -> (index)
+// USE-EXPENSIVE-MATH: %{{.*}} = affine.apply #[[$MAP_APPLY]](%[[IV_MAIN]])[%[[BOUND]]]
+// USE-EXPENSIVE-MATH: %[[TAIL_RES:.*]] = affine.for %[[IV_TAIL:.*]] = 4 to #[[$MAP_UB]]()[%[[BOUND]]] iter_args(%[[ARG_TAIL:.*]] = %[[MAIN_RES]]) -> (index)
+// USE-EXPENSIVE-MATH: %{{.*}} = affine.apply #[[$MAP_APPLY]](%[[IV_TAIL]])[%[[BOUND]]]
+
+// -----
+
+// USE-EXPENSIVE-MATH-DAG: #[[$MAP_APPLY:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
+// USE-EXPENSIVE-MATH-LABEL: func @fully_constantized_no_peeling
+func.func @fully_constantized_no_peeling() {
+ %c0 = arith.constant 0 : index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ affine.for %iv = %bound to 6 step 2 iter_args(%arg = %c0) -> index {
+ %sum = arith.addi %arg, %bound : index
+ affine.yield %sum : index
+ }
+ return
+}
+
+// USE-EXPENSIVE-MATH: %[[C0:.*]] = arith.constant 0 : index
+// USE-EXPENSIVE-MATH: %[[BOUND:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
+// USE-EXPENSIVE-MATH: %{{.*}} = affine.for %[[IV:.*]] = 0 to 3 iter_args(%{{.*}} = %[[C0]]) -> (index)
+// USE-EXPENSIVE-MATH: %{{.*}} = affine.apply #[[$MAP_APPLY]](%[[IV]])[%[[BOUND]]]
+
+// -----
+
+// USE-EXPENSIVE-MATH-DAG: #[[$MAP_REMAIN:.*]] = affine_map<()[s0] -> ((-s0 + 7) ceildiv 2)>
+// USE-EXPENSIVE-MATH-LABEL: func @peel_nested_loops
+func.func @peel_nested_loops() {
+ %c0 = arith.constant 0 : index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ affine.for %i = %bound to 7 step 2 {
+ affine.for %j = %bound to 7 step 2 {
+ "test.foo"() : () -> ()
+ }
+ }
+ return
+}
+
+// USE-EXPENSIVE-MATH: %[[BOUND:.*]] = test.value_with_bounds
+// USE-EXPENSIVE-MATH: affine.for %{{.*}} = 0 to 3
+// USE-EXPENSIVE-MATH: affine.for %{{.*}} = 0 to 3
+// USE-EXPENSIVE-MATH: affine.for %{{.*}} = 3 to #[[$MAP_REMAIN]]()[%[[BOUND]]]
+// USE-EXPENSIVE-MATH: affine.for %{{.*}} = 3 to #[[$MAP_REMAIN]]()[%[[BOUND]]]
+// USE-EXPENSIVE-MATH: affine.for %{{.*}} = 0 to 3
+// USE-EXPENSIVE-MATH: affine.for %{{.*}} = 3 to #[[$MAP_REMAIN]]()[%[[BOUND]]]
+
+// -----
+
+// USE-EXPENSIVE-MATH-AND-PROMOTE-LABEL: func @single_iter_promoted
+func.func @single_iter_promoted() {
+ %c0 = arith.constant 0 :index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ affine.for %iv = %bound to 2 step 2 {
+ "test.foo"() : () -> ()
+ }
+ return
+}
+
+// USE-EXPENSIVE-MATH-AND-PROMOTE-NOT: affine.for
+
+// -----
+
+// USE-EXPENSIVE-MATH-AND-PROMOTE-LABEL: func @single_iter_promoted_with_remainder
+func.func @single_iter_promoted_with_remainder() {
+ %c0 = arith.constant 0 : index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ affine.for %i = %bound to 3 step 2 {
+ "test.foo"() : () -> ()
+ }
+ return
+}
+
+// USE-EXPENSIVE-MATH-AND-PROMOTE: "test.foo"() : () -> ()
+// USE-EXPENSIVE-MATH-AND-PROMOTE: affine.for
+// USE-EXPENSIVE-MATH-AND-PROMOTE: "test.foo"() : () -> ()
More information about the Mlir-commits
mailing list