[Mlir-commits] [mlir] [mlir][affine] Add useExpensiveMath option to AffineLoopNormalize pass (PR #211989)
lonely eagle
llvmlistbot at llvm.org
Tue Jul 28 19:58:23 PDT 2026
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/211989
>From 559a1cebe5d06c1b21e84ea5f40cfbebca7a219d Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sat, 25 Jul 2026 05:56:24 +0000
Subject: [PATCH 1/4] add affine-loop-constantize-bounds pass.
---
.../mlir/Dialect/Affine/Transforms/Passes.td | 13 +++
.../AffineLoopConstantizeBounds.cpp | 104 ++++++++++++++++++
.../Dialect/Affine/Transforms/CMakeLists.txt | 1 +
.../affine-loop-constantize-bounds.mlir | 45 ++++++++
4 files changed, 163 insertions(+)
create mode 100644 mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
create mode 100644 mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
index 7cecbcac5b932..44cb733c525b9 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
@@ -465,4 +465,17 @@ def AffineFoldMemRefAliasOps : Pass<"affine-fold-memref-alias-ops"> {
let dependentDialects = ["memref::MemRefDialect"];
}
+def AffineLoopConstantizeBounds : Pass<"affine-loop-constantize-bounds"> {
+ let summary = "Constantize affine loop bounds using range analysis";
+ let description = [{
+ The pass leverages value range analysis (e.g., via the Presburger library)
+ to transform affine loops with dynamic symbolic bounds into a static-trip-count
+ main loop and a residual tail loop.
+
+ By rewriting the main loop to have strict compile-time constant bounds
+ (typically normalized to `0` to `TC_safe` with step `1`), it exposes the loop
+ body to downstream affine optimization pipelines.
+ }];
+}
+
#endif // MLIR_DIALECT_AFFINE_PASSES
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
new file mode 100644
index 0000000000000..529d4043f5693
--- /dev/null
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
@@ -0,0 +1,104 @@
+//===- AffineLoopConstantizeBounds.cpp - Constantize loop bounds pass ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to transform affine loops with symbolic bounds
+// into a static-trip-count main loop and a residual tail loop using range
+// analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Transforms/Passes.h"
+#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include <cstdint>
+
+namespace mlir {
+namespace affine {
+#define GEN_PASS_DEF_AFFINELOOPCONSTANTIZEBOUNDS
+#include "mlir/Dialect/Affine/Transforms/Passes.h.inc"
+} // namespace affine
+} // namespace mlir
+
+#define DEBUG_TYPE "affine-loop-constantize-bounds"
+
+using namespace mlir;
+using namespace mlir::affine;
+
+namespace {
+
+/// Computes the constant upper or lower bound for a given affine map expression
+/// \p map and its operands \p operands, constrained by the specified \p type.
+static FailureOr<int64_t> computeConstantBound(AffineMap map,
+ ValueRange operands,
+ presburger::BoundType type) {
+ ValueBoundsConstraintSet::Variable var(map, operands);
+ return ValueBoundsConstraintSet::computeConstantBound(type, var, nullptr,
+ {true, true});
+}
+
+static LogicalResult inferAffineLoopUpperConstantBound(AffineForOp forOp) {
+ // Ensure the loop is normalized (lower bound is strictly 0 and step is 1).
+ if (!forOp.hasConstantLowerBound() || forOp.getConstantLowerBound() != 0)
+ return failure();
+ if (forOp.getStepAsInt() != 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();
+
+ IRRewriter b(forOp->getContext());
+
+ // 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());
+ 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);
+ return success();
+}
+
+struct AffineLoopConstantizeBounds
+ : public affine::impl::AffineLoopConstantizeBoundsBase<
+ AffineLoopConstantizeBounds> {
+ void runOnOperation() override;
+};
+} // namespace
+
+void AffineLoopConstantizeBounds::runOnOperation() {
+ SmallVector<AffineForOp, 4> loops;
+ getOperation()->walk([&](AffineForOp forOp) {
+ // First collect eligible loops and normalize them so that their lower bound
+ // is locked to 0 and step to 1.
+ if (succeeded(normalizeAffineFor(forOp, false)))
+ loops.push_back(forOp);
+ });
+ // Infer and rewrite the upper bound into a compile-time constant for each
+ // loop.
+ for (AffineForOp loop : loops)
+ (void)inferAffineLoopUpperConstantBound(loop);
+}
diff --git a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
index 9d912139810b2..3a854e1c7dc36 100644
--- a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRAffineTransforms
AffineDataCopyGeneration.cpp
AffineExpandIndexOps.cpp
AffineExpandIndexOpsAsAffine.cpp
+ AffineLoopConstantizeBounds.cpp
AffineLoopInvariantCodeMotion.cpp
AffineLoopNormalize.cpp
AffineParallelize.cpp
diff --git a/mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir b/mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir
new file mode 100644
index 0000000000000..6bb9aed223e83
--- /dev/null
+++ b/mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir
@@ -0,0 +1,45 @@
+// RUN: mlir-opt %s -affine-loop-constantize-bounds -split-input-file | FileCheck %s
+
+// CHECK-DAG: #[[$MAP_APPLY:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
+// CHECK-DAG: #[[$MAP_UB:.+]] = affine_map<()[s0] -> ((-s0 + 9) ceildiv 2)>
+
+// CHECK-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
+}
+
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[BOUND:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
+// CHECK: %[[MAIN_RES:.*]] = affine.for %[[IV_MAIN:.*]] = 0 to 4 iter_args(%[[ARG_MAIN:.*]] = %[[C0]]) -> (index) {
+// CHECK: %{{.*}} = affine.apply #map(%[[IV_MAIN]])[%[[BOUND]]]
+// CHECK: }
+// CHECK: %[[TAIL_RES:.*]] = affine.for %[[IV_TAIL:.*]] = 4 to #map1()[%[[BOUND]]] iter_args(%[[ARG_TAIL:.*]] = %[[MAIN_RES]]) -> (index) {
+// CHECK: %{{.*}} = affine.apply #map(%[[IV_TAIL]])[%[[BOUND]]]
+// CHECK: }
+
+// -----
+// CHECK: #map = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
+// CHECK-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
+}
+
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[BOUND:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
+// CHECK: %{{.*}} = affine.for %[[IV:.*]] = 0 to 3 iter_args(%{{.*}} = %[[C0]]) -> (index) {
+// CHECK: %{{.*}} = affine.apply #map(%[[IV]])[%[[BOUND]]]
+// CHECK: }
>From dbe3e12c7d80edb8585bfc6794ebaa204a622569 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 27 Jul 2026 08:45:50 +0000
Subject: [PATCH 2/4] fix nit.
---
.../AffineLoopConstantizeBounds.cpp | 23 +++++++++++++------
1 file changed, 16 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
index 529d4043f5693..ed424d2638631 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
@@ -15,6 +15,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Transforms/Passes.h"
#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include <cstdint>
@@ -33,16 +34,24 @@ using namespace mlir::affine;
namespace {
/// Computes the constant upper or lower bound for a given affine map expression
-/// \p map and its operands \p operands, constrained by the specified \p type.
+/// 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,
- {true, true});
+ return ValueBoundsConstraintSet::computeConstantBound(
+ type, var, nullptr, {/*closedUb=*/true, /*allowIntegerType=*/true});
}
-static LogicalResult inferAffineLoopUpperConstantBound(AffineForOp forOp) {
+/// 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(AffineForOp forOp,
+ RewriterBase &b) {
// Ensure the loop is normalized (lower bound is strictly 0 and step is 1).
if (!forOp.hasConstantLowerBound() || forOp.getConstantLowerBound() != 0)
return failure();
@@ -61,8 +70,6 @@ static LogicalResult inferAffineLoopUpperConstantBound(AffineForOp forOp) {
if (failed(upperMin) || *upperMin <= 0)
return failure();
- IRRewriter b(forOp->getContext());
-
// 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).
@@ -97,8 +104,10 @@ void AffineLoopConstantizeBounds::runOnOperation() {
if (succeeded(normalizeAffineFor(forOp, false)))
loops.push_back(forOp);
});
+
// Infer and rewrite the upper bound into a compile-time constant for each
// loop.
+ IRRewriter b(&getContext());
for (AffineForOp loop : loops)
- (void)inferAffineLoopUpperConstantBound(loop);
+ (void)inferAffineLoopUpperConstantBound(loop, b);
}
>From c29d0ec2538d02e24ccf3888140e16bec65d70fc Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 28 Jul 2026 12:01:11 +0000
Subject: [PATCH 3/4] update code.
---
.../mlir/Dialect/Affine/Transforms/Passes.h | 3 +-
.../mlir/Dialect/Affine/Transforms/Passes.td | 15 +--
.../AffineLoopConstantizeBounds.cpp | 113 ------------------
.../Affine/Transforms/AffineLoopNormalize.cpp | 85 ++++++++++++-
.../Dialect/Affine/Transforms/CMakeLists.txt | 1 -
.../affine-loop-constantize-bounds.mlir | 45 -------
.../Dialect/Affine/affine-loop-normalize.mlir | 62 ++++++++++
7 files changed, 148 insertions(+), 176 deletions(-)
delete mode 100644 mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
delete mode 100644 mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir
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 44cb733c525b9..1f8e57f776607 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 the Presburger to constantize loop upper bounds (higher analysis cost)">
];
}
@@ -465,17 +467,4 @@ def AffineFoldMemRefAliasOps : Pass<"affine-fold-memref-alias-ops"> {
let dependentDialects = ["memref::MemRefDialect"];
}
-def AffineLoopConstantizeBounds : Pass<"affine-loop-constantize-bounds"> {
- let summary = "Constantize affine loop bounds using range analysis";
- let description = [{
- The pass leverages value range analysis (e.g., via the Presburger library)
- to transform affine loops with dynamic symbolic bounds into a static-trip-count
- main loop and a residual tail loop.
-
- By rewriting the main loop to have strict compile-time constant bounds
- (typically normalized to `0` to `TC_safe` with step `1`), it exposes the loop
- body to downstream affine optimization pipelines.
- }];
-}
-
#endif // MLIR_DIALECT_AFFINE_PASSES
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
deleted file mode 100644
index ed424d2638631..0000000000000
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
+++ /dev/null
@@ -1,113 +0,0 @@
-//===- AffineLoopConstantizeBounds.cpp - Constantize loop bounds pass ---===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-//
-// This file implements a pass to transform affine loops with symbolic bounds
-// into a static-trip-count main loop and a residual tail loop using range
-// analysis.
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Affine/Transforms/Passes.h"
-#include "mlir/Dialect/Affine/Utils.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Interfaces/ValueBoundsOpInterface.h"
-#include <cstdint>
-
-namespace mlir {
-namespace affine {
-#define GEN_PASS_DEF_AFFINELOOPCONSTANTIZEBOUNDS
-#include "mlir/Dialect/Affine/Transforms/Passes.h.inc"
-} // namespace affine
-} // namespace mlir
-
-#define DEBUG_TYPE "affine-loop-constantize-bounds"
-
-using namespace mlir;
-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(AffineForOp forOp,
- RewriterBase &b) {
- // Ensure the loop is normalized (lower bound is strictly 0 and step is 1).
- if (!forOp.hasConstantLowerBound() || forOp.getConstantLowerBound() != 0)
- return failure();
- if (forOp.getStepAsInt() != 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());
- 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);
- return success();
-}
-
-struct AffineLoopConstantizeBounds
- : public affine::impl::AffineLoopConstantizeBoundsBase<
- AffineLoopConstantizeBounds> {
- void runOnOperation() override;
-};
-} // namespace
-
-void AffineLoopConstantizeBounds::runOnOperation() {
- SmallVector<AffineForOp, 4> loops;
- getOperation()->walk([&](AffineForOp forOp) {
- // First collect eligible loops and normalize them so that their lower bound
- // is locked to 0 and step to 1.
- if (succeeded(normalizeAffineFor(forOp, false)))
- loops.push_back(forOp);
- });
-
- // Infer and rewrite the upper bound into a compile-time constant for each
- // loop.
- IRRewriter b(&getContext());
- for (AffineForOp loop : loops)
- (void)inferAffineLoopUpperConstantBound(loop, b);
-}
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
index c74d147016237..bf03dfcef7e1e 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
@@ -10,11 +10,13 @@
//
//===----------------------------------------------------------------------===//
+#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"
namespace mlir {
namespace affine {
@@ -28,13 +30,79 @@ 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) {
+ // Ensure the loop is normalized (lower bound is strictly 0 and step is 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(forOp);
+
+ 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 +112,23 @@ 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());
+ getOperation()->walk([&](AffineForOp forOp) {
+ (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/lib/Dialect/Affine/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
index 3a854e1c7dc36..9d912139810b2 100644
--- a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
@@ -2,7 +2,6 @@ add_mlir_dialect_library(MLIRAffineTransforms
AffineDataCopyGeneration.cpp
AffineExpandIndexOps.cpp
AffineExpandIndexOpsAsAffine.cpp
- AffineLoopConstantizeBounds.cpp
AffineLoopInvariantCodeMotion.cpp
AffineLoopNormalize.cpp
AffineParallelize.cpp
diff --git a/mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir b/mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir
deleted file mode 100644
index 6bb9aed223e83..0000000000000
--- a/mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir
+++ /dev/null
@@ -1,45 +0,0 @@
-// RUN: mlir-opt %s -affine-loop-constantize-bounds -split-input-file | FileCheck %s
-
-// CHECK-DAG: #[[$MAP_APPLY:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
-// CHECK-DAG: #[[$MAP_UB:.+]] = affine_map<()[s0] -> ((-s0 + 9) ceildiv 2)>
-
-// CHECK-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
-}
-
-// CHECK: %[[C0:.*]] = arith.constant 0 : index
-// CHECK: %[[BOUND:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
-// CHECK: %[[MAIN_RES:.*]] = affine.for %[[IV_MAIN:.*]] = 0 to 4 iter_args(%[[ARG_MAIN:.*]] = %[[C0]]) -> (index) {
-// CHECK: %{{.*}} = affine.apply #map(%[[IV_MAIN]])[%[[BOUND]]]
-// CHECK: }
-// CHECK: %[[TAIL_RES:.*]] = affine.for %[[IV_TAIL:.*]] = 4 to #map1()[%[[BOUND]]] iter_args(%[[ARG_TAIL:.*]] = %[[MAIN_RES]]) -> (index) {
-// CHECK: %{{.*}} = affine.apply #map(%[[IV_TAIL]])[%[[BOUND]]]
-// CHECK: }
-
-// -----
-// CHECK: #map = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
-// CHECK-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
-}
-
-// CHECK: %[[C0:.*]] = arith.constant 0 : index
-// CHECK: %[[BOUND:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
-// CHECK: %{{.*}} = affine.for %[[IV:.*]] = 0 to 3 iter_args(%{{.*}} = %[[C0]]) -> (index) {
-// CHECK: %{{.*}} = affine.apply #map(%[[IV]])[%[[BOUND]]]
-// CHECK: }
diff --git a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
index 7d90efec0c21b..1db1d11a92edd 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,63 @@ 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: }
+// 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: }
+
+// -----
+
+// 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: }
+
+// -----
+
+// USE-EXPENSIVE-MATH-AND-PROMOTE-LABEL: func @constantized_single_iter_promoted
+
+func.func @constantized_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 iter_args(%arg = %c0) -> index {
+ %sum = arith.addi %arg, %bound : index
+ affine.yield %sum : index
+ }
+ return
+}
+
+// USE-EXPENSIVE-MATH-AND-PROMOTE-NOT: affine.for
>From 5b99371bfcfc367abab70a2469b7be52f546d922 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Wed, 29 Jul 2026 02:58:05 +0000
Subject: [PATCH 4/4] update code.
---
.../Affine/Transforms/AffineLoopNormalize.cpp | 16 +++++++++++++---
1 file changed, 13 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
index bf03dfcef7e1e..fd7b500ab301e 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
@@ -17,6 +17,7 @@
#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 {
@@ -50,7 +51,8 @@ static FailureOr<int64_t> computeConstantBound(AffineMap map,
static LogicalResult
inferAffineLoopUpperConstantBound(RewriterBase &b, AffineForOp forOp,
bool promoteSingleIter = true) {
- // Ensure the loop is normalized (lower bound is strictly 0 and step is 1).
+ // 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)
@@ -117,9 +119,17 @@ struct AffineLoopNormalizePass
// loop.
if (useExpensiveMath) {
IRRewriter b(&getContext());
- getOperation()->walk([&](AffineForOp forOp) {
+ 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);
- });
}
}
};
More information about the Mlir-commits
mailing list