[Mlir-commits] [mlir] 0390898 - [mlir][affine] Implement LoopLikeInterface::getStaticTripCount on AffineForOp (#204687)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 19 07:11:54 PDT 2026
Author: Jeremy Kun
Date: 2026-06-19T07:11:48-07:00
New Revision: 0390898335f9f32ea71ff288a5b4085cecc10391
URL: https://github.com/llvm/llvm-project/commit/0390898335f9f32ea71ff288a5b4085cecc10391
DIFF: https://github.com/llvm/llvm-project/commit/0390898335f9f32ea71ff288a5b4085cecc10391.diff
LOG: [mlir][affine] Implement LoopLikeInterface::getStaticTripCount on AffineForOp (#204687)
LoopLikeInterface is useful, but missing `getStaticTripCount` requires
adding extra cases to check when processing otherwise dialect-agnostic
code.
There is an existing free function `getConstantTripCount`, which I
deprecated and replaced (NFC) with the new implementation. I believe the
new implementation is slightly more efficient than
`getConstantTripCount` because it checks if the expression is constant
and fast-fails before constructing the output `AffineMap` that was
returned by `getTripCountMapAndOperands`.
Assisted by Gemini
Added:
mlir/test/Dialect/Affine/trip-count.mlir
Modified:
mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
mlir/lib/Dialect/Affine/Analysis/Utils.cpp
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp
mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
index 43d61832cafdd..3fcb63a4da885 100644
--- a/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
+++ b/mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h
@@ -41,6 +41,7 @@ void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *map,
/// Returns the trip count of the loop if it's a constant, std::nullopt
/// otherwise. This uses affine expression analysis and is able to determine
/// constant trip count in non-trivial cases.
+[[deprecated("use AffineForOp::getStaticTripCount instead")]]
std::optional<uint64_t> getConstantTripCount(AffineForOp forOp);
/// Returns the greatest known integral divisor of the trip count. Affine
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 3d7cbcc375d2a..1e14f9f37288d 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -135,7 +135,7 @@ def AffineForOp : Affine_Op<"for",
RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getLoopInductionVars", "getLoopLowerBounds", "getLoopSteps",
"getLoopUpperBounds", "getYieldedValuesMutable",
- "replaceWithAdditionalYields"]>,
+ "replaceWithAdditionalYields", "getStaticTripCount"]>,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
["getEntrySuccessorOperands", "getSuccessorInputs"]>]> {
let summary = "for operation";
diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index 166d39e88d41e..40802cc6e85e5 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -214,27 +214,9 @@ void mlir::affine::getTripCountMapAndOperands(
/// getTripCount) and is able to determine constant trip count in non-trivial
/// cases.
std::optional<uint64_t> mlir::affine::getConstantTripCount(AffineForOp forOp) {
- SmallVector<Value, 4> operands;
- AffineMap map;
- getTripCountMapAndOperands(forOp, &map, &operands);
-
- if (!map)
- return std::nullopt;
-
- // Take the min if all trip counts are constant.
- std::optional<uint64_t> tripCount;
- for (auto resultExpr : map.getResults()) {
- if (auto constExpr = dyn_cast<AffineConstantExpr>(resultExpr)) {
- if (tripCount.has_value())
- tripCount =
- std::min(*tripCount, static_cast<uint64_t>(constExpr.getValue()));
- else
- tripCount = constExpr.getValue();
- } else {
- return std::nullopt;
- }
- }
- return tripCount;
+ if (std::optional<APInt> tripCount = forOp.getStaticTripCount())
+ return tripCount->getZExtValue();
+ return std::nullopt;
}
/// Returns the greatest known integral divisor of the trip count. Affine
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index ebe932a14694a..cac305df8ba75 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -1833,9 +1833,9 @@ bool mlir::affine::buildSliceTripCountMap(
forOp.getConstantUpperBound() - forOp.getConstantLowerBound();
continue;
}
- std::optional<uint64_t> maybeConstTripCount = getConstantTripCount(forOp);
+ std::optional<APInt> maybeConstTripCount = forOp.getStaticTripCount();
if (maybeConstTripCount.has_value()) {
- (*tripCountMap)[op] = *maybeConstTripCount;
+ (*tripCountMap)[op] = maybeConstTripCount->getZExtValue();
continue;
}
return false;
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 7d8974bd6c1b7..f095500495f18 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2826,6 +2826,61 @@ std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopUpperBounds() {
OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))};
}
+std::optional<APInt> AffineForOp::getStaticTripCount() {
+ MLIRContext *context = getContext();
+ int64_t step = getStepAsInt();
+ if (step <= 0)
+ return std::nullopt;
+
+ if (hasConstantBounds()) {
+ int64_t lb = getConstantLowerBound();
+ int64_t ub = getConstantUpperBound();
+ int64_t loopSpan = ub - lb;
+ if (loopSpan < 0)
+ loopSpan = 0;
+ return APInt(64, llvm::divideCeilSigned(loopSpan, step));
+ }
+
+ auto lbMap = getLowerBoundMap();
+ auto ubMap = getUpperBoundMap();
+ if (lbMap.getNumResults() != 1)
+ return std::nullopt;
+
+ // Difference of each upper bound expression from the single lower bound
+ // expression (divided by the step) provides the expressions for the trip
+ // count map.
+ AffineValueMap ubValueMap(ubMap, getUpperBoundOperands());
+
+ SmallVector<AffineExpr, 4> lbSplatExpr(ubValueMap.getNumResults(),
+ lbMap.getResult(0));
+ auto lbMapSplat = AffineMap::get(lbMap.getNumDims(), lbMap.getNumSymbols(),
+ lbSplatExpr, context);
+ AffineValueMap lbSplatValueMap(lbMapSplat, getLowerBoundOperands());
+
+ AffineValueMap tripCountValueMap;
+ AffineValueMap::
diff erence(ubValueMap, lbSplatValueMap, &tripCountValueMap);
+
+ // Take the min if all trip counts are constant.
+ std::optional<uint64_t> tripCount;
+ for (unsigned i = 0, e = tripCountValueMap.getNumResults(); i < e; ++i) {
+ AffineExpr expr = tripCountValueMap.getResult(i).ceilDiv(step);
+ if (auto constExpr = llvm::dyn_cast<AffineConstantExpr>(expr)) {
+ uint64_t value = constExpr.getValue();
+ if (tripCount.has_value())
+ tripCount = std::min(*tripCount, value);
+ else
+ tripCount = value;
+ } else {
+ return std::nullopt;
+ }
+ }
+
+ if (tripCount.has_value())
+ return APInt(64, *tripCount);
+
+ return std::nullopt;
+}
+
FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(
RewriterBase &rewriter, ValueRange newInitOperands,
bool replaceInitOperandUsesInLoop,
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
index 3c55830df61c3..1887c321e206a 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
@@ -178,8 +178,9 @@ void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
// at least once. For unknown (dynamic) or zero trip counts we cannot prove
// the body executes, so hoisting a side-effectful op would change observable
// program semantics. Pure (side-effect-free) ops may always be hoisted.
- auto tripCount = getConstantTripCount(forOp);
- bool guaranteedToExecute = tripCount.has_value() && *tripCount > 0;
+ auto tripCount = forOp.getStaticTripCount();
+ bool guaranteedToExecute =
+ tripCount.has_value() && tripCount->getZExtValue() > 0;
for (Operation &op : *forOp.getBody()) {
// Register op in the set of ops that have users. This set is used
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
index 188db218a5220..d3208d5c8f7eb 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
@@ -91,12 +91,13 @@ static void adjustToDivisorsOfTripCounts(ArrayRef<AffineForOp> band,
assert(band.size() == tileSizes->size() && "invalid tile size count");
for (unsigned i = 0, e = band.size(); i < e; i++) {
unsigned &tSizeAdjusted = (*tileSizes)[i];
- std::optional<uint64_t> mayConst = getConstantTripCount(band[i]);
+ AffineForOp forOp = band[i];
+ std::optional<APInt> mayConst = forOp.getStaticTripCount();
if (!mayConst)
continue;
// Adjust the tile size to largest factor of the trip count less than
// tSize.
- uint64_t constTripCount = *mayConst;
+ uint64_t constTripCount = mayConst->getZExtValue();
if (constTripCount > 1 && tSizeAdjusted > constTripCount / 2)
tSizeAdjusted = constTripCount / 2;
while (constTripCount % tSizeAdjusted != 0)
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp
index 837d4f714d25e..1006a7d2c3cca 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp
@@ -100,8 +100,8 @@ void LoopUnroll::runOnOperation() {
// so that loops are gathered from innermost to outermost (or else
// unrolling an outer one may delete gathered inner ones).
getOperation().walk([&](AffineForOp forOp) {
- std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
- if (tripCount && *tripCount <= unrollFullThreshold)
+ std::optional<APInt> tripCount = forOp.getStaticTripCount();
+ if (tripCount && tripCount->getZExtValue() <= unrollFullThreshold)
loops.push_back(forOp);
});
for (auto forOp : loops)
diff --git a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
index d84cb4f0cde5f..575b529658127 100644
--- a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
@@ -245,8 +245,7 @@ static void findMatchingStartFinishInsts(
/// 'forOp' is deleted, and a prologue, a new pipelined loop, and epilogue are
/// inserted right before where it was.
void PipelineDataTransfer::runOnAffineForOp(AffineForOp forOp) {
- auto mayBeConstTripCount = getConstantTripCount(forOp);
- if (!mayBeConstTripCount) {
+ if (!forOp.getStaticTripCount()) {
LLVM_DEBUG(forOp.emitRemark("won't pipeline due to unknown trip count"));
return;
}
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 82247dcfe71ef..68296ea3368a1 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -357,7 +357,7 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
bool siblingFusionUser) {
// Check if the reduction loop is a single iteration loop.
- std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
+ std::optional<APInt> tripCount = forOp.getStaticTripCount();
if (!tripCount || *tripCount != 1)
return failure();
auto *parentOp = forOp->getParentOp();
@@ -496,14 +496,14 @@ bool mlir::affine::getLoopNestStats(AffineForOp forOpRoot,
// Record trip count for 'forOp'. Set flag if trip count is not
// constant.
- std::optional<uint64_t> maybeConstTripCount = getConstantTripCount(forOp);
+ std::optional<APInt> maybeConstTripCount = forOp.getStaticTripCount();
if (!maybeConstTripCount) {
// Currently only constant trip count loop nests are supported.
LDBG() << "Non-constant trip count unsupported";
return WalkResult::interrupt();
}
- stats->tripCountMap[childForOp] = *maybeConstTripCount;
+ stats->tripCountMap[childForOp] = maybeConstTripCount->getZExtValue();
return WalkResult::advance();
});
return !walkResult.wasInterrupted();
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 8f1249e3afaf0..90bc57e950cf1 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -117,7 +117,7 @@ static void replaceIterArgsAndYieldResults(AffineForOp forOp) {
/// 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);
+ std::optional<APInt> tripCount = forOp.getStaticTripCount();
if (!tripCount || *tripCount != 1)
return failure();
@@ -239,12 +239,12 @@ LogicalResult mlir::affine::affineForOpBodySkew(AffineForOp forOp,
// conditional guards (or context information to prevent such versioning). The
// better way to pipeline for such loops is to first tile them and extract
// constant trip count "full tiles" before applying this.
- auto mayBeConstTripCount = getConstantTripCount(forOp);
+ auto mayBeConstTripCount = forOp.getStaticTripCount();
if (!mayBeConstTripCount) {
LLVM_DEBUG(forOp.emitRemark("non-constant trip count loop not handled"));
return success();
}
- uint64_t tripCount = *mayBeConstTripCount;
+ uint64_t tripCount = mayBeConstTripCount->getZExtValue();
assert(isOpwiseShiftValid(forOp, shifts) &&
"shifts will lead to an invalid transformation\n");
@@ -707,8 +707,10 @@ constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
// Bounds for intra-tile loops.
for (unsigned i = 0; i < width; i++) {
int64_t largestDiv = getLargestDivisorOfTripCount(origLoops[i]);
- std::optional<uint64_t> mayBeConstantCount =
- getConstantTripCount(origLoops[i]);
+ AffineForOp forOp = origLoops[i];
+ std::optional<uint64_t> mayBeConstantCount = std::nullopt;
+ if (auto staticTripCount = forOp.getStaticTripCount())
+ mayBeConstantCount = staticTripCount->getZExtValue();
// The lower bound is just the tile-space loop.
AffineMap lbMap = b.getDimIdentityMap();
newLoops[width + i].setLowerBound(
@@ -869,9 +871,9 @@ void mlir::affine::getPerfectlyNestedLoops(
/// Unrolls this loop completely.
LogicalResult mlir::affine::loopUnrollFull(AffineForOp forOp) {
- std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
+ std::optional<APInt> mayBeConstantTripCount = forOp.getStaticTripCount();
if (mayBeConstantTripCount.has_value()) {
- uint64_t tripCount = *mayBeConstantTripCount;
+ uint64_t tripCount = mayBeConstantTripCount->getZExtValue();
if (tripCount == 0)
return success();
if (tripCount == 1)
@@ -885,10 +887,10 @@ LogicalResult mlir::affine::loopUnrollFull(AffineForOp forOp) {
/// whichever is lower.
LogicalResult mlir::affine::loopUnrollUpToFactor(AffineForOp forOp,
uint64_t unrollFactor) {
- std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
+ std::optional<APInt> mayBeConstantTripCount = forOp.getStaticTripCount();
if (mayBeConstantTripCount.has_value() &&
- *mayBeConstantTripCount < unrollFactor)
- return loopUnrollByFactor(forOp, *mayBeConstantTripCount);
+ mayBeConstantTripCount->ult(unrollFactor))
+ return loopUnrollByFactor(forOp, mayBeConstantTripCount->getZExtValue());
return loopUnrollByFactor(forOp, unrollFactor);
}
@@ -998,7 +1000,9 @@ LogicalResult mlir::affine::loopUnrollByFactor(
bool cleanUpUnroll) {
assert(unrollFactor > 0 && "unroll factor should be positive");
- std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
+ std::optional<uint64_t> mayBeConstantTripCount = std::nullopt;
+ if (auto staticTripCount = forOp.getStaticTripCount())
+ mayBeConstantTripCount = staticTripCount->getZExtValue();
if (unrollFactor == 1) {
if (mayBeConstantTripCount == 1 && failed(promoteIfSingleIteration(forOp)))
return failure();
@@ -1060,10 +1064,10 @@ LogicalResult mlir::affine::loopUnrollByFactor(
LogicalResult mlir::affine::loopUnrollJamUpToFactor(AffineForOp forOp,
uint64_t unrollJamFactor) {
- std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
+ std::optional<APInt> mayBeConstantTripCount = forOp.getStaticTripCount();
if (mayBeConstantTripCount.has_value() &&
- *mayBeConstantTripCount < unrollJamFactor)
- return loopUnrollJamByFactor(forOp, *mayBeConstantTripCount);
+ mayBeConstantTripCount->getZExtValue() < unrollJamFactor)
+ return loopUnrollJamByFactor(forOp, mayBeConstantTripCount->getZExtValue());
return loopUnrollJamByFactor(forOp, unrollJamFactor);
}
@@ -1085,7 +1089,9 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
uint64_t unrollJamFactor) {
assert(unrollJamFactor > 0 && "unroll jam factor should be positive");
- std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
+ std::optional<uint64_t> mayBeConstantTripCount = std::nullopt;
+ if (auto staticTripCount = forOp.getStaticTripCount())
+ mayBeConstantTripCount = staticTripCount->getZExtValue();
if (unrollJamFactor == 1) {
if (mayBeConstantTripCount == 1 && failed(promoteIfSingleIteration(forOp)))
return failure();
diff --git a/mlir/test/Dialect/Affine/trip-count.mlir b/mlir/test/Dialect/Affine/trip-count.mlir
new file mode 100644
index 0000000000000..e28e410fd2112
--- /dev/null
+++ b/mlir/test/Dialect/Affine/trip-count.mlir
@@ -0,0 +1,38 @@
+// This test ensures that the LoopLikeInterfaceOp methods required
+// for op-agnostic trip count analysis work for affine.for.
+
+// RUN: mlir-opt %s -test-scf-for-utils --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @affine_constant_loops
+func.func @affine_constant_loops() {
+ // CHECK: "test.trip-count" = 10
+ affine.for %i = 0 to 10 {
+ affine.yield
+ }
+ // CHECK: "test.trip-count" = 5
+ affine.for %i = 0 to 10 step 2 {
+ affine.yield
+ }
+ // CHECK: "test.trip-count" = 0
+ affine.for %i = 10 to 0 {
+ affine.yield
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @affine_symbolic_loops
+func.func @affine_symbolic_loops(%N : index) {
+ // CHECK: "test.trip-count" = "none"
+ affine.for %i = 0 to %N {
+ affine.yield
+ }
+
+ // CHECK: "test.trip-count" = 4
+ affine.for %i = max affine_map<(d0) -> (d0)>(%N) to min affine_map<(d0) -> (d0 + 4)>(%N) {
+ affine.yield
+ }
+
+ return
+}
More information about the Mlir-commits
mailing list