[Mlir-commits] [mlir] [mlir][affine] Use value bound inference to determine minimum/maximum trip counts in loop analysis (PR #128113)
lonely eagle
llvmlistbot at llvm.org
Tue May 26 01:26:04 PDT 2026
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/128113
>From f8bb05d98116d9800bf22ef1e895a164ba99b346 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 25 May 2026 03:12:54 +0000
Subject: [PATCH 1/5] update code.
---
.../Dialect/Affine/Analysis/LoopAnalysis.h | 4 ++
.../Dialect/Affine/Analysis/LoopAnalysis.cpp | 59 +++++++++++++++++++
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 6 ++
3 files changed, 69 insertions(+)
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
index 43d61832cafdd..1f5b017e5b189 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
@@ -15,7 +15,9 @@
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h"
+#include <cstdint>
#include <optional>
+#include <utility>
namespace mlir {
class AffineExpr;
@@ -43,6 +45,8 @@ void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *map,
/// constant trip count in non-trivial cases.
std::optional<uint64_t> getConstantTripCount(AffineForOp forOp);
+std::optional<std::pair<int64_t, int64_t>> getTripCount(AffineForOp forOp);
+
/// Returns the greatest known integral divisor of the trip count. Affine
/// expression analysis is used (indirectly through getTripCount), and
/// this method is thus able to determine non-trivial divisors.
diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index 166d39e88d41e..9a3d45b74b910 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -12,18 +12,27 @@
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
+#include "mlir/Analysis/Presburger/PresburgerRelation.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
#include "mlir/Dialect/Affine/Analysis/NestedMatcher.h"
#include "mlir/Dialect/Affine/Analysis/Utils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
+#include <cstdint>
#include <numeric>
#include <optional>
+#include <utility>
#define DEBUG_TYPE "affine-loop-analysis"
@@ -237,6 +246,56 @@ std::optional<uint64_t> mlir::affine::getConstantTripCount(AffineForOp forOp) {
return tripCount;
}
+static std::optional<uint64_t>
+getKnownTripCountBound(AffineMap map, SmallVectorImpl<Value> &operands,
+ presburger::BoundType type) {
+ std::optional<uint64_t> tripCount;
+ for (unsigned i = 0, e = map.getResults().size(); i < e; ++i) {
+ AffineMap subMap = map.getSubMap(i);
+ ValueBoundsConstraintSet::Variable var(subMap, operands);
+ auto lbBound = ValueBoundsConstraintSet::computeConstantBound(
+ mlir::presburger::BoundType::LB, var);
+ auto ubBound = ValueBoundsConstraintSet::computeConstantBound(
+ mlir::presburger::BoundType::UB, var, nullptr);
+ if (failed(lbBound) || failed(ubBound))
+ return std::nullopt;
+ if (type == presburger::BoundType::LB) {
+ if (tripCount.has_value())
+ tripCount =
+ std::min(*tripCount, static_cast<uint64_t>(lbBound.value()));
+ else
+ tripCount = lbBound.value();
+ } else if (type == presburger::BoundType::UB) {
+ if (tripCount.has_value())
+ tripCount =
+ std::max(*tripCount, static_cast<uint64_t>(ubBound.value()));
+ else
+ tripCount = ubBound.value();
+ } else {
+ return std::nullopt;
+ }
+ }
+ return tripCount;
+}
+
+std::optional<std::pair<int64_t, int64_t>>
+mlir::affine::getTripCount(AffineForOp forOp) {
+ SmallVector<Value, 4> operands;
+ AffineMap map;
+ getTripCountMapAndOperands(forOp, &map, &operands);
+
+ if (!map)
+ return {};
+
+ std::optional<int64_t> minTrip =
+ getKnownTripCountBound(map, operands, presburger::BoundType::LB);
+ std::optional<int64_t> maxTrip =
+ getKnownTripCountBound(map, operands, presburger::BoundType::UB);
+ if (!minTrip || !maxTrip)
+ return {};
+ return std::make_pair(*minTrip, *maxTrip);
+}
+
/// Returns the greatest known integral divisor of the trip count. Affine
/// expression analysis is used (indirectly through getTripCount), and
/// this method is thus able to determine non-trivial divisors.
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 8f1249e3afaf0..29bb531ec3288 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -27,7 +27,9 @@
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/raw_ostream.h"
+#include <cstdint>
#include <optional>
+#include <utility>
#define DEBUG_TYPE "loop-utils"
@@ -878,6 +880,10 @@ LogicalResult mlir::affine::loopUnrollFull(AffineForOp forOp) {
return promoteIfSingleIteration(forOp);
return loopUnrollByFactor(forOp, tripCount);
}
+ std::optional<std::pair<int64_t, int64_t>> tripCountRange =
+ getTripCount(forOp);
+ if (tripCountRange.has_value()) {
+ }
return failure();
}
>From a2a6e9cc0a6db06268738866b483f84405861c20 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 26 May 2026 03:27:04 +0000
Subject: [PATCH 2/5] update code.
---
.../Dialect/Affine/Analysis/LoopAnalysis.h | 3 +-
.../Dialect/Affine/Analysis/LoopAnalysis.cpp | 43 ++++++--
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 97 ++++++++++++++++++-
3 files changed, 127 insertions(+), 16 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
index 1f5b017e5b189..7cc249ecd99c0 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
@@ -45,7 +45,8 @@ void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *map,
/// constant trip count in non-trivial cases.
std::optional<uint64_t> getConstantTripCount(AffineForOp forOp);
-std::optional<std::pair<int64_t, int64_t>> getTripCount(AffineForOp forOp);
+std::optional<std::pair<uint64_t, uint64_t>>
+computeLoopTripCountConstantBounds(AffineForOp forOp);
/// Returns the greatest known integral divisor of the trip count. Affine
/// expression analysis is used (indirectly through getTripCount), and
diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index 9a3d45b74b910..048b13699b647 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -23,8 +23,11 @@
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Value.h"
+#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Support/Debug.h"
@@ -246,17 +249,27 @@ std::optional<uint64_t> mlir::affine::getConstantTripCount(AffineForOp forOp) {
return tripCount;
}
+static FailureOr<int64_t> computeConstantBound(AffineMap map,
+ ValueRange operands,
+ presburger::BoundType type) {
+ ValueBoundsConstraintSet::Variable var(map, operands);
+ ValueBoundsOptions options;
+ options.closedUB = true;
+ options.allowIntegerType = true;
+ return ValueBoundsConstraintSet::computeConstantBound(type, var, nullptr,
+ options);
+}
+
static std::optional<uint64_t>
getKnownTripCountBound(AffineMap map, SmallVectorImpl<Value> &operands,
presburger::BoundType type) {
std::optional<uint64_t> tripCount;
for (unsigned i = 0, e = map.getResults().size(); i < e; ++i) {
AffineMap subMap = map.getSubMap(i);
- ValueBoundsConstraintSet::Variable var(subMap, operands);
- auto lbBound = ValueBoundsConstraintSet::computeConstantBound(
- mlir::presburger::BoundType::LB, var);
- auto ubBound = ValueBoundsConstraintSet::computeConstantBound(
- mlir::presburger::BoundType::UB, var, nullptr);
+ FailureOr<int64_t> lbBound =
+ computeConstantBound(subMap, operands, presburger::BoundType::LB);
+ FailureOr<int64_t> ubBound =
+ computeConstantBound(subMap, operands, presburger::BoundType::UB);
if (failed(lbBound) || failed(ubBound))
return std::nullopt;
if (type == presburger::BoundType::LB) {
@@ -278,8 +291,8 @@ getKnownTripCountBound(AffineMap map, SmallVectorImpl<Value> &operands,
return tripCount;
}
-std::optional<std::pair<int64_t, int64_t>>
-mlir::affine::getTripCount(AffineForOp forOp) {
+std::optional<std::pair<uint64_t, uint64_t>>
+mlir::affine::computeLoopTripCountConstantBounds(AffineForOp forOp) {
SmallVector<Value, 4> operands;
AffineMap map;
getTripCountMapAndOperands(forOp, &map, &operands);
@@ -287,9 +300,9 @@ mlir::affine::getTripCount(AffineForOp forOp) {
if (!map)
return {};
- std::optional<int64_t> minTrip =
+ std::optional<uint64_t> minTrip =
getKnownTripCountBound(map, operands, presburger::BoundType::LB);
- std::optional<int64_t> maxTrip =
+ std::optional<uint64_t> maxTrip =
getKnownTripCountBound(map, operands, presburger::BoundType::UB);
if (!minTrip || !maxTrip)
return {};
@@ -311,7 +324,7 @@ uint64_t mlir::affine::getLargestDivisorOfTripCount(AffineForOp forOp) {
// divisors.
assert(map.getNumResults() >= 1 && "expected one or more results");
std::optional<uint64_t> gcd;
- for (auto resultExpr : map.getResults()) {
+ for (auto [idx, resultExpr] : llvm::enumerate(map.getResults())) {
uint64_t thisGcd;
if (auto constExpr = dyn_cast<AffineConstantExpr>(resultExpr)) {
uint64_t tripCount = constExpr.getValue();
@@ -321,6 +334,16 @@ uint64_t mlir::affine::getLargestDivisorOfTripCount(AffineForOp forOp) {
else
// The greatest divisor is the trip count.
thisGcd = tripCount;
+ } else if (FailureOr<int64_t> lbBound = computeConstantBound(
+ map.getSubMap(idx), operands, presburger::BoundType::LB),
+ ubBound = computeConstantBound(map.getSubMap(idx), operands,
+ presburger::BoundType::UB);
+ !failed(lbBound) && !failed(ubBound)) {
+ if (*lbBound == *ubBound)
+ thisGcd =
+ (*lbBound == 0) ? std::numeric_limits<uint64_t>::max() : *lbBound;
+ else
+ thisGcd = 1;
} else {
// Trip count is not a known constant; return its largest known divisor.
thisGcd = resultExpr.getLargestKnownDivisor();
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 29bb531ec3288..2eb66108bdae6 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -22,10 +22,12 @@
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/IntegerSet.h"
#include "mlir/IR/OperationSupport.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
+#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/raw_ostream.h"
#include <cstdint>
#include <optional>
@@ -116,12 +118,73 @@ static void replaceIterArgsAndYieldResults(AffineForOp forOp) {
std::get<0>(e).replaceAllUsesWith(std::get<1>(e));
}
+/// Return true if we can prove that the we always run at least the first
+/// iteration of the ForOp.
+static bool alwaysRunsFirstIteration(AffineForOp op) {
+ // Can't perform the analysis if the loops's bounds aren't index-typed.
+ if (!op.getInductionVar().getType().isIndex())
+ return false;
+ if (op.getLowerBoundMap().getNumResults() != 1 ||
+ op.getUpperBoundMap().getNumResults() != 1)
+ return false;
+
+ SmallVector<Value> lowerMapOperands = op.getLowerBoundOperands();
+ SmallVector<Value> upperMapOperands = op.getUpperBoundOperands();
+ ValueBoundsConstraintSet::Variable lower(op.getLowerBoundMap(),
+ lowerMapOperands);
+ ValueBoundsConstraintSet::Variable upper(op.getUpperBoundMap(),
+ upperMapOperands);
+ FailureOr<bool> isLb = ValueBoundsConstraintSet::compare(
+ lower, ValueBoundsConstraintSet::LT, upper);
+ return isLb.value_or(false);
+}
+
+/// Return true if we can prove that the we never run more than one iteration of
+/// the ForOp.
+static bool neverRunsSecondIteration(AffineForOp op) {
+ // Can't perform the analysis if the loops's bounds aren't index-typed.
+ if (!op.getInductionVar().getType().isIndex())
+ return false;
+
+ if (op.getLowerBoundMap().getNumResults() != 1 ||
+ op.getUpperBoundMap().getNumResults() != 1)
+ return false;
+
+ // The loop will only loop once if the inducation variable for the next time
+ // in the loop is greater than or equal to upper.
+ MLIRContext *context = op.getContext();
+ SmallVector<Value> lowerMapOperands = op.getLowerBoundOperands();
+ SmallVector<Value> upperMapOperands = op.getUpperBoundOperands();
+ SmallVector<AffineExpr> results;
+ AffineMap lowerMap = op.getLowerBoundMap();
+ for (AffineExpr expr : lowerMap.getResults())
+ results.push_back(expr + op.getStep().getSExtValue());
+
+ AffineMap nextItMap = AffineMap::get(
+ lowerMap.getNumDims(), lowerMap.getNumSymbols(), results, context);
+ ValueBoundsConstraintSet::Variable nextItVar(nextItMap, lowerMapOperands);
+ ValueBoundsConstraintSet::Variable upperVar(op.getUpperBoundMap(),
+ upperMapOperands);
+ FailureOr<bool> isUpperUnderNextIter = ValueBoundsConstraintSet::compare(
+ nextItVar, ValueBoundsConstraintSet::GE, upperVar);
+ return isUpperUnderNextIter.value_or(false);
+}
+
/// Promotes the loop body of a forOp to its containing block if the forOp
/// was known to have a single iteration.
LogicalResult mlir::affine::promoteIfSingleIteration(AffineForOp forOp) {
std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
- if (!tripCount || *tripCount != 1)
+ // Only allow loops that are guaranteed to execute exactly once. If the trip
+ // count is constant, it must be exactly. If the trip count is dynamic, verify
+ // via affine analysis that it always runs the first iteration but never
+ // reaches the second.
+ if (tripCount && *tripCount != 1) {
+ return failure();
+ }
+ if (!tripCount &&
+ !(alwaysRunsFirstIteration(forOp) && neverRunsSecondIteration(forOp))) {
return failure();
+ }
// TODO: extend this for arbitrary affine bounds.
if (forOp.getLowerBoundMap().getNumResults() != 1)
@@ -881,8 +944,15 @@ LogicalResult mlir::affine::loopUnrollFull(AffineForOp forOp) {
return loopUnrollByFactor(forOp, tripCount);
}
std::optional<std::pair<int64_t, int64_t>> tripCountRange =
- getTripCount(forOp);
- if (tripCountRange.has_value()) {
+ computeLoopTripCountConstantBounds(forOp);
+ if (tripCountRange) {
+ int64_t minTrip = tripCountRange->first;
+ int64_t maxTrip = tripCountRange->second;
+ if (minTrip == 0)
+ return success();
+ if (minTrip == maxTrip && minTrip == 1)
+ return promoteIfSingleIteration(forOp);
+ return loopUnrollByFactor(forOp, minTrip);
}
return failure();
}
@@ -1005,8 +1075,22 @@ LogicalResult mlir::affine::loopUnrollByFactor(
assert(unrollFactor > 0 && "unroll factor should be positive");
std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
+ std::optional<uint64_t> mayBeConstantMaxTripCount = mayBeConstantTripCount;
+ bool tripEqual = true;
+ if (!mayBeConstantTripCount.has_value()) {
+ std::optional<std::pair<uint64_t, uint64_t>> tripBound =
+ computeLoopTripCountConstantBounds(forOp);
+ if (tripBound.has_value()) {
+ mayBeConstantTripCount = tripBound->first;
+ if (tripBound->first != tripBound->second) {
+ mayBeConstantMaxTripCount = tripBound->second;
+ tripEqual = tripBound->first == tripBound->second;
+ }
+ }
+ }
if (unrollFactor == 1) {
- if (mayBeConstantTripCount == 1 && failed(promoteIfSingleIteration(forOp)))
+ if (mayBeConstantTripCount == 1 && tripEqual &&
+ failed(promoteIfSingleIteration(forOp)))
return failure();
return success();
}
@@ -1016,7 +1100,10 @@ LogicalResult mlir::affine::loopUnrollByFactor(
return success();
// If the trip count is lower than the unroll factor, no unrolled body.
- if (mayBeConstantTripCount && *mayBeConstantTripCount < unrollFactor) {
+ if ((mayBeConstantTripCount && tripEqual &&
+ *mayBeConstantTripCount < unrollFactor) ||
+ (mayBeConstantMaxTripCount &&
+ *mayBeConstantMaxTripCount < unrollFactor)) {
if (cleanUpUnroll) {
// Unroll the cleanup loop if cleanUpUnroll is specified.
return loopUnrollFull(forOp);
>From b9c48d702aed587114df4ccef3d42696c8fc8b73 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 26 May 2026 06:39:25 +0000
Subject: [PATCH 3/5] update code.
---
.../Dialect/Affine/Analysis/LoopAnalysis.cpp | 43 +++----------------
1 file changed, 7 insertions(+), 36 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index 048b13699b647..0a4a14cadd1b0 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -260,37 +260,6 @@ static FailureOr<int64_t> computeConstantBound(AffineMap map,
options);
}
-static std::optional<uint64_t>
-getKnownTripCountBound(AffineMap map, SmallVectorImpl<Value> &operands,
- presburger::BoundType type) {
- std::optional<uint64_t> tripCount;
- for (unsigned i = 0, e = map.getResults().size(); i < e; ++i) {
- AffineMap subMap = map.getSubMap(i);
- FailureOr<int64_t> lbBound =
- computeConstantBound(subMap, operands, presburger::BoundType::LB);
- FailureOr<int64_t> ubBound =
- computeConstantBound(subMap, operands, presburger::BoundType::UB);
- if (failed(lbBound) || failed(ubBound))
- return std::nullopt;
- if (type == presburger::BoundType::LB) {
- if (tripCount.has_value())
- tripCount =
- std::min(*tripCount, static_cast<uint64_t>(lbBound.value()));
- else
- tripCount = lbBound.value();
- } else if (type == presburger::BoundType::UB) {
- if (tripCount.has_value())
- tripCount =
- std::max(*tripCount, static_cast<uint64_t>(ubBound.value()));
- else
- tripCount = ubBound.value();
- } else {
- return std::nullopt;
- }
- }
- return tripCount;
-}
-
std::optional<std::pair<uint64_t, uint64_t>>
mlir::affine::computeLoopTripCountConstantBounds(AffineForOp forOp) {
SmallVector<Value, 4> operands;
@@ -299,12 +268,14 @@ mlir::affine::computeLoopTripCountConstantBounds(AffineForOp forOp) {
if (!map)
return {};
+ if (map.getNumResults() != 1)
+ return {};
- std::optional<uint64_t> minTrip =
- getKnownTripCountBound(map, operands, presburger::BoundType::LB);
- std::optional<uint64_t> maxTrip =
- getKnownTripCountBound(map, operands, presburger::BoundType::UB);
- if (!minTrip || !maxTrip)
+ FailureOr<int64_t> minTrip =
+ computeConstantBound(map, operands, presburger::BoundType::LB);
+ FailureOr<int64_t> maxTrip =
+ computeConstantBound(map, operands, presburger::BoundType::UB);
+ if (failed(minTrip) || failed(maxTrip))
return {};
return std::make_pair(*minTrip, *maxTrip);
}
>From 02473852701da060aa4b28b1afb6974f4367c0cc Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 26 May 2026 07:41:31 +0000
Subject: [PATCH 4/5] add comment and add test.
---
.../Dialect/Affine/Analysis/LoopAnalysis.h | 2 -
.../Dialect/Affine/Analysis/LoopAnalysis.cpp | 25 +++----
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 12 ++-
mlir/test/Dialect/Affine/unroll.mlir | 75 +++++++++++++++++++
4 files changed, 96 insertions(+), 18 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
index 7cc249ecd99c0..f8cb28cf8dd69 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
@@ -15,9 +15,7 @@
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h"
-#include <cstdint>
#include <optional>
-#include <utility>
namespace mlir {
class AffineExpr;
diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index 0a4a14cadd1b0..319d4c2d6f74d 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -12,30 +12,19 @@
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
-#include "mlir/Analysis/Presburger/IntegerRelation.h"
-#include "mlir/Analysis/Presburger/PresburgerRelation.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
#include "mlir/Dialect/Affine/Analysis/NestedMatcher.h"
#include "mlir/Dialect/Affine/Analysis/Utils.h"
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
-#include "mlir/IR/AffineMap.h"
-#include "mlir/IR/Value.h"
-#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
-#include <cstdint>
#include <numeric>
#include <optional>
-#include <utility>
#define DEBUG_TYPE "affine-loop-analysis"
@@ -260,6 +249,9 @@ static FailureOr<int64_t> computeConstantBound(AffineMap map,
options);
}
+/// Computes the constant lower and upper bounds of the given affine loop's
+/// trip count using the Presburger-based `ValueBoundsConstraintSet`
+/// infrastructure.
std::optional<std::pair<uint64_t, uint64_t>>
mlir::affine::computeLoopTripCountConstantBounds(AffineForOp forOp) {
SmallVector<Value, 4> operands;
@@ -268,6 +260,7 @@ mlir::affine::computeLoopTripCountConstantBounds(AffineForOp forOp) {
if (!map)
return {};
+ // Currently, we only support trip count maps with a single result expression.
if (map.getNumResults() != 1)
return {};
@@ -310,11 +303,17 @@ uint64_t mlir::affine::getLargestDivisorOfTripCount(AffineForOp forOp) {
ubBound = computeConstantBound(map.getSubMap(idx), operands,
presburger::BoundType::UB);
!failed(lbBound) && !failed(ubBound)) {
- if (*lbBound == *ubBound)
+ if (*lbBound == *ubBound) {
+ // If the sub-map yields a strict constant bound (LB == UB), this
+ // specific dimension's trip count is completely static.
thisGcd =
(*lbBound == 0) ? std::numeric_limits<uint64_t>::max() : *lbBound;
- else
+ } else {
+ // By fundamental number theory, any two consecutive integers are
+ // coprime (gcd(n, n+1) = 1). Therefore, the overall alignment factor
+ // for this dynamic range strictly collapses to 1.
thisGcd = 1;
+ }
} else {
// Trip count is not a known constant; return its largest known divisor.
thisGcd = resultExpr.getLargestKnownDivisor();
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 2eb66108bdae6..e19b88492b7e6 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -27,11 +27,8 @@
#include "llvm/ADT/MapVector.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
-#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/raw_ostream.h"
-#include <cstdint>
#include <optional>
-#include <utility>
#define DEBUG_TYPE "loop-utils"
@@ -1074,15 +1071,24 @@ LogicalResult mlir::affine::loopUnrollByFactor(
bool cleanUpUnroll) {
assert(unrollFactor > 0 && "unroll factor should be positive");
+ // Attempt to resolve the loop trip count as a literal constant first.
std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
std::optional<uint64_t> mayBeConstantMaxTripCount = mayBeConstantTripCount;
+
+ // Optimistically assume the loop has a fixed/constant trip count.
bool tripEqual = true;
if (!mayBeConstantTripCount.has_value()) {
+ // Fall back to the polyhedral bounds infrastructure using
+ // `ValueBoundsConstraintSet`.
std::optional<std::pair<uint64_t, uint64_t>> tripBound =
computeLoopTripCountConstantBounds(forOp);
if (tripBound.has_value()) {
+ // Treat the lower bound as the baseline trip count, if first == second,
+ // the bounds have collapsed into a constant.
mayBeConstantTripCount = tripBound->first;
if (tripBound->first != tripBound->second) {
+ // Guaranteed to evaluate to 'false' here, signaling a true dynamic
+ // interval [LB, UB].
mayBeConstantMaxTripCount = tripBound->second;
tripEqual = tripBound->first == tripBound->second;
}
diff --git a/mlir/test/Dialect/Affine/unroll.mlir b/mlir/test/Dialect/Affine/unroll.mlir
index efdceed7c9a25..5da786364560f 100644
--- a/mlir/test/Dialect/Affine/unroll.mlir
+++ b/mlir/test/Dialect/Affine/unroll.mlir
@@ -12,6 +12,7 @@
// UNROLL-FULL-DAG: [[$MAP4:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 + 1)>
// UNROLL-FULL-DAG: [[$MAP5:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 + 3)>
// UNROLL-FULL-DAG: [[$MAP6:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 + s0 + 1)>
+// UNROLL-FULL-DAG: [[$MAP7:#map[0-9]*]] = affine_map<()[s0] -> (s0 + (((-s0 + 9) ceildiv 2) floordiv 4) * 8)>
// SHORT-DAG: [[$MAP0:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
@@ -23,6 +24,7 @@
// UNROLL-BY-4-DAG: [[$MAP5:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 + s0 + 1)>
// UNROLL-BY-4-DAG: [[$MAP6:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 * 16 + d1)>
// UNROLL-BY-4-DAG: [[$MAP11:#map[0-9]*]] = affine_map<(d0) -> (d0)>
+// UNROLL-BY-4-DAG: [[$MAP8:#map[0-9]*]] = affine_map<()[s0] -> (s0 + (((-s0 + 11) ceildiv 2) floordiv 4) * 8)>
// UNROLL-FULL-LABEL: func @loop_nest_simplest() {
func.func @loop_nest_simplest() {
@@ -258,6 +260,58 @@ gpu.module @unroll_full {
}
}
+// UNROLL-FULL-LABEL: func @bound_unroll_partial
+func.func @bound_unroll_partial() {
+ %c0 = arith.constant 0 :index
+ // UNROLL-FULL: %[[C0:.*]] = arith.constant 0 : index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ affine.for %iv = %bound to 3 step 2 iter_args(%arg = %c0) -> index {
+ %sum = arith.addi %arg, %c0 : index
+ affine.yield %sum : index
+ }
+ // UNROLL-FULL: affine.for %{{.*}} = %{{.*}} to 3 step 2 iter_args(%[[ARG:.*]] = %[[C0]]) -> (index) {
+ // UNROLL-FULL-NEXT: %[[SUM:.*]] = arith.addi %[[ARG]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: affine.yield %[[SUM]] : index
+ // UNROLL-FULL-NEXT: }
+ return
+}
+
+// UNROLL-FULL-LABEL: func @bound_unroll_all
+func.func @bound_unroll_all() {
+ %c0 = arith.constant 0 :index
+ // UNROLL-FULL: %[[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, %c0 : index
+ affine.yield %sum : index
+ }
+ // UNROLL-FULL: %[[SUM_0:.*]] = arith.addi %[[C0]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: %[[SUM_1:.*]] = arith.addi %[[SUM_0]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: %[[SUM_2:.*]] = arith.addi %[[SUM_1]], %[[C0]] : index
+ return
+}
+
+// UNROLL-FULL-LABEL: func.func @bound_partial_unroll_factor_4
+func.func @bound_partial_unroll_factor_4() {
+ %c0 = arith.constant 0 :index
+ // UNROLL-FULL: %[[C0:.*]] = arith.constant 0 : index
+ // UNROLL-FULL: %[[Bound:.*]] = test.value_with_bounds {max = 1 : index, min = 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, %c0 : index
+ affine.yield %sum : index
+ }
+ // UNROLL-FULL-NEXT: %[[SUM_0:.*]] = arith.addi %[[C0]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: %[[SUM_1:.*]] = arith.addi %[[SUM_0]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: %[[SUM_2:.*]] = arith.addi %[[SUM_1]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: %[[SUM_3:.*]] = arith.addi %[[SUM_2]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: affine.for %{{.*}} = [[$MAP7]]()[%[[Bound]]] to 9 step 2 iter_args(%[[ARG:.*]] = %[[SUM_3]]) -> (index) {
+ // UNROLL-FULL-NEXT: %[[SUM_4:.*]] = arith.addi %[[ARG]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: affine.yield %[[SUM_4]] : index
+ // UNROLL-FULL-NEXT: }
+ return
+}
+
// SHORT-LABEL: func @loop_nest_outer_unroll() {
func.func @loop_nest_outer_unroll() {
// SHORT: affine.for %arg0 = 0 to 4 {
@@ -701,6 +755,27 @@ func.func @unroll_with_iter_args_and_promotion(%arg0 : f32, %arg1 : f32) -> f32
return %sum : f32
}
+// UNROLL-BY-4-LABEL: func @bound_unroll_by_factor_4
+func.func @bound_unroll_by_factor_4() {
+ %c0 = arith.constant 0 :index
+ // UNROLL-BY-4: %[[C0:.*]] = arith.constant 0 : index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ // UNROLL-BY-4: %[[Bound:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
+ affine.for %iv = %bound to 11 step 2 iter_args(%arg = %c0) -> index {
+ %sum = arith.addi %arg, %c0 : index
+ affine.yield %sum : index
+ }
+ // UNROLL-BY-4-NEXT: %[[SUM_0:.*]] = arith.addi %[[C0]], %[[C0]] : index
+ // UNROLL-BY-4-NEXT: %[[SUM_1:.*]] = arith.addi %[[SUM_0]], %[[C0]] : index
+ // UNROLL-BY-4-NEXT: %[[SUM_2:.*]] = arith.addi %[[SUM_1]], %[[C0]] : index
+ // UNROLL-BY-4-NEXT: %[[SUM_3:.*]] = arith.addi %[[SUM_2]], %[[C0]] : index
+ // UNROLL-BY-4-NEXT: affine.for %[[VAL_20:.*]] = [[$MAP8]](){{\[}}%[[Bound]]] to 11 step 2 iter_args(%[[ARG:.*]] = %[[SUM_3]]) -> (index) {
+ // UNROLL-BY-4-NEXT: %[[SUM_4:.*]] = arith.addi %[[ARG]], %[[C0]] : index
+ // UNROLL-BY-4-NEXT: affine.yield %[[SUM_4]] : index
+ // UNROLL-BY-4-NEXT: }
+ return
+}
+
// UNROLL-FULL: func @unroll_zero_trip_count_case
func.func @unroll_zero_trip_count_case() {
// CHECK-NEXT: affine.for %{{.*}} = 0 to 0
>From dbca15b9f23fdea1409ccb0d114549b98e306b13 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 26 May 2026 08:24:11 +0000
Subject: [PATCH 5/5] add comment and add test.
---
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 5 +----
mlir/test/Dialect/Affine/unroll.mlir | 24 +++++++++++++++++++++
2 files changed, 25 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index e19b88492b7e6..a648c68649440 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1106,10 +1106,7 @@ LogicalResult mlir::affine::loopUnrollByFactor(
return success();
// If the trip count is lower than the unroll factor, no unrolled body.
- if ((mayBeConstantTripCount && tripEqual &&
- *mayBeConstantTripCount < unrollFactor) ||
- (mayBeConstantMaxTripCount &&
- *mayBeConstantMaxTripCount < unrollFactor)) {
+ if (mayBeConstantTripCount && *mayBeConstantTripCount < unrollFactor) {
if (cleanUpUnroll) {
// Unroll the cleanup loop if cleanUpUnroll is specified.
return loopUnrollFull(forOp);
diff --git a/mlir/test/Dialect/Affine/unroll.mlir b/mlir/test/Dialect/Affine/unroll.mlir
index 5da786364560f..d263c104e2564 100644
--- a/mlir/test/Dialect/Affine/unroll.mlir
+++ b/mlir/test/Dialect/Affine/unroll.mlir
@@ -13,6 +13,7 @@
// UNROLL-FULL-DAG: [[$MAP5:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 + 3)>
// UNROLL-FULL-DAG: [[$MAP6:#map[0-9]*]] = affine_map<(d0)[s0] -> (d0 + s0 + 1)>
// UNROLL-FULL-DAG: [[$MAP7:#map[0-9]*]] = affine_map<()[s0] -> (s0 + (((-s0 + 9) ceildiv 2) floordiv 4) * 8)>
+// UNROLL-FULL-DAG: [[$MAP8:#map[0-9]*]] = affine_map<()[s0, s1] -> (s0 + (((-s0 + s1) ceildiv 2) floordiv
// SHORT-DAG: [[$MAP0:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
@@ -312,6 +313,29 @@ func.func @bound_partial_unroll_factor_4() {
return
}
+// UNROLL-FULL-LABEL: func @bound_for_unroll_dynamic_lower_and_upper
+func.func @bound_for_unroll_dynamic_lower_and_upper() {
+ %c0 = arith.constant 0 :index
+ %lower_bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ %upper_bound = test.value_with_bounds { min = 6 : index, max = 7 : index}
+ affine.for %iv = %lower_bound to %upper_bound step 2 iter_args(%arg = %c0) -> index {
+ %sum = arith.addi %arg, %c0 : index
+ affine.yield %sum : index
+ }
+ // UNROLL-FULL-NEXT: %[[C0:.*]] = arith.constant 0 : index
+ // UNROLL-FULL-NEXT: %[[LB:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
+ // UNROLL-FULL-NEXT: %[[UB:.*]] = test.value_with_bounds {max = 7 : index, min = 6 : index}
+ // UNROLL-FULL-NEXT: %[[ADDI_0:.*]] = arith.addi %[[C0]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: %[[ADDI_1:.*]] = arith.addi %[[ADDI_0]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: %[[UNROLLED_INIT:.*]] = arith.addi %[[ADDI_1]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: %[[LOOP_RES:.*]] = affine.for %{{.*}} = [[$MAP8]]()
+ // UNROLL-FULL-SAME: [%[[LB]], %[[UB]]] to %[[UB]] step 2 iter_args(%[[VAL:.*]] = %[[UNROLLED_INIT]]) -> (index) {
+ // UNROLL-FULL-NEXT: %[[YIELD_VAL:.*]] = arith.addi %[[VAL]], %[[C0]] : index
+ // UNROLL-FULL-NEXT: affine.yield %[[YIELD_VAL]] : index
+ // UNROLL-FULL-NEXT: }
+ return
+}
+
// SHORT-LABEL: func @loop_nest_outer_unroll() {
func.func @loop_nest_outer_unroll() {
// SHORT: affine.for %arg0 = 0 to 4 {
More information about the Mlir-commits
mailing list