[Mlir-commits] [mlir] [mlir][affine] Update getSliceBounds to allow multi-result upper bound maps (PR #219369)
Uday Bondhugula
llvmlistbot at llvm.org
Sun Aug 30 02:41:30 PDT 2026
https://github.com/bondhugula updated https://github.com/llvm/llvm-project/pull/219369
>From 607372da4d6a1a1929b1abfa049030514a08bcaf Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday at polymagelabs.com>
Date: Wed, 26 Aug 2026 16:41:24 +0530
Subject: [PATCH] [mlir][affine] Update getSliceBounds to allow multi-result
upper bound maps
Resolve a long-standing TODO on supporting multi-result upper bound maps
on affine analysis utility getSliceBounds. This makes affine fusion more
powerful.
`getSliceBounds` threw away any upper bound that came out of more than
one inequality and put the constant bound in its place, under a TODO
saying it was conservative until `getConstDifference` in LoopFusion could
handle multiple bounds (b/126426796). Several inequalities is what a
destination loop clamped at the end of the data produces -- of a 1000
long dim tiled by 64, the region a tile reads is bounded by
`min(%i * 64 + 64, 1000)` -- and dropping that for the constant leaves
`1000`, which says only that the slice ends somewhere before the end of
the data. A slice of one 64-wide tile then costs as a slice of everything
from the tile onwards, and fusion refuses it as redundant computation
that isn't there.
Keep such a bound. It is already the min of its results, which is the
form an affine.for upper bound takes, so materializing the slice needs
nothing new. `getConstDifference` takes each result against the lower
bound and keeps the smallest difference that comes out constant: every
result bounds the count from above, so the smallest constant among them
is the tightest constant bound there is. On the example above the
tile-relative result gives 64 and the extent gives nothing constant.
Only a caller that can consume such a bound should be handed one, so this
is behind `allowMultiResultUb`, which the two slice-computation callers
pass and the value-bounds one does not -- it can't express a bound of
several results and would lose the bounds it computes today.
Assisted-by: Claude.
---
.../Analysis/FlatLinearValueConstraints.h | 14 +++---
.../Analysis/FlatLinearValueConstraints.cpp | 17 ++++---
mlir/lib/Dialect/Affine/Analysis/Utils.cpp | 44 +++++++++++++------
.../Affine/loop-fusion-slice-computation.mlir | 29 ++++++++++++
4 files changed, 79 insertions(+), 25 deletions(-)
diff --git a/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h b/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
index 79fb201935f0d..23c7418fd13a2 100644
--- a/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
+++ b/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
@@ -10,9 +10,9 @@
#define MLIR_ANALYSIS_FLATLINEARVALUECONSTRAINTS_H
#include "mlir/Analysis/Presburger/IntegerRelation.h"
-#include "mlir/Analysis/Presburger/Matrix.h"
#include "mlir/IR/AffineExpr.h"
-#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/ValueRange.h"
#include <optional>
namespace mlir {
@@ -20,7 +20,6 @@ namespace mlir {
class AffineMap;
class IntegerSet;
class MLIRContext;
-class Value;
class MemRefType;
struct MutableAffineMap;
@@ -158,10 +157,15 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
///
/// By default the returned lower bounds are closed and upper bounds are open.
/// If `closedUb` is true, the upper bound is closed.
+ ///
+ /// An upper bound built from more than one inequality is the min of them.
+ /// Only a caller that can consume such a bound should ask for it, via
+ /// `allowMultiResultUb`; the rest are given the constant upper bound
+ /// instead, which is weaker but always a single result.
void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context,
SmallVectorImpl<AffineMap> *lbMaps,
- SmallVectorImpl<AffineMap> *ubMaps,
- bool closedUB = false);
+ SmallVectorImpl<AffineMap> *ubMaps, bool closedUB = false,
+ bool allowMultiResultUB = false);
/// Composes an affine map whose dimensions and symbols match one to one with
/// the dimensions and symbols of this FlatLinearConstraints. The results of
diff --git a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
index 6d28387b0e5c8..79367fc857379 100644
--- a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
+++ b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
@@ -8,8 +8,8 @@
#include "mlir/Analysis//FlatLinearValueConstraints.h"
+#include "mlir/Analysis/Presburger/PWMAFunction.h"
#include "mlir/Analysis/Presburger/PresburgerSpace.h"
-#include "mlir/Analysis/Presburger/Simplex.h"
#include "mlir/Analysis/Presburger/Utils.h"
#include "mlir/IR/AffineExprVisitor.h"
#include "mlir/IR/Builders.h"
@@ -701,7 +701,8 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
MLIRContext *context,
SmallVectorImpl<AffineMap> *lbMaps,
SmallVectorImpl<AffineMap> *ubMaps,
- bool closedUB) {
+ bool closedUB,
+ bool allowMultiResultUB) {
assert(offset + num <= getNumDimVars() && "invalid range");
// Basic simplification.
@@ -753,9 +754,6 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
// If the above fails, we'll just use the constant lower bound and the
// constant upper bound (if they exist) as the slice bounds.
- // TODO: being conservative for the moment in cases that
- // lead to multiple bounds - until getConstDifference in LoopFusion.cpp is
- // fixed (b/126426796).
if (!lbMap || lbMap.getNumResults() != 1) {
LLVM_DEBUG(llvm::dbgs()
<< "WARNING: Potentially over-approximating slice lb\n");
@@ -765,7 +763,14 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
getAffineConstantExpr(*lbConst, context));
}
}
- if (!ubMap || ubMap.getNumResults() != 1) {
+ // An upper bound of several results is the min of them, for e.g.
+ // the bound of a tiled loop clamped at the end of the data looks like:
+ // `min(%i * 64 + 64, 1000)`. Falling back to the constant bound
+ // the way the lower one does would drop the tile-relative result and
+ // leave nothing but the extent of the data, turning a slice of a single
+ // tile into a slice of everything from the tile onwards.
+ if (!ubMap || ubMap.getNumResults() == 0 ||
+ (!allowMultiResultUB && ubMap.getNumResults() != 1)) {
LLVM_DEBUG(llvm::dbgs()
<< "WARNING: Potentially over-approximating slice ub\n");
auto ubConst = getConstantBound64(BoundType::UB, pos + offset);
diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 321c8e34d907c..17c629c5c1d11 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -1604,8 +1604,7 @@ unsigned mlir::affine::getInnermostCommonLoopDepth(
unsigned loopDepthLimit = std::numeric_limits<unsigned>::max();
for (unsigned i = 0; i < numOps; ++i) {
getAffineForIVs(*ops[i], &loops[i]);
- loopDepthLimit =
- std::min(loopDepthLimit, static_cast<unsigned>(loops[i].size()));
+ loopDepthLimit = std::min(loopDepthLimit, (unsigned)loops[i].size());
}
unsigned loopDepth = 0;
@@ -1754,7 +1753,8 @@ mlir::affine::computeSliceUnion(ArrayRef<Operation *> opsA,
// Get slice bounds from slice union constraints 'sliceUnionCst'.
sliceUnionCst.getSliceBounds(/*offset=*/0, numSliceLoopIVs,
opsA[0]->getContext(), &sliceUnion->lbs,
- &sliceUnion->ubs);
+ &sliceUnion->ubs, /*closedUb=*/false,
+ /*allowMultiResultUb=*/true);
// Add slice bound operands of union.
SmallVector<Value, 4> sliceBoundOperands;
@@ -1791,21 +1791,36 @@ mlir::affine::computeSliceUnion(ArrayRef<Operation *> opsA,
return SliceComputationResult::Success;
}
-// TODO: extend this to handle multiple result maps.
+/// Returns the number of iterations the slice bounded below by `lbMap` and
+/// above by `ubMap` runs for, where that is a constant.
+///
+/// An upper bound of several results is the min of them, so each result taken
+/// against the lower bound bounds the count from above and the smallest of
+/// those that comes out constant is the tightest constant bound there is. A
+/// tiled loop clamped at the end of the data has exactly this shape --
+/// `min(%i * 64 + 64, 1000)` over `%i * 64` -- where the tile-relative result
+/// gives the 64 and the extent gives nothing constant at all.
static std::optional<uint64_t> getConstDifference(AffineMap lbMap,
AffineMap ubMap) {
- assert(lbMap.getNumResults() == 1 && "expected single result bound map");
- assert(ubMap.getNumResults() == 1 && "expected single result bound map");
+ assert(lbMap.getNumResults() == 1 && "expected single result lower bound");
+ assert(ubMap.getNumResults() >= 1 && "expected at least one upper bound");
assert(lbMap.getNumDims() == ubMap.getNumDims());
assert(lbMap.getNumSymbols() == ubMap.getNumSymbols());
AffineExpr lbExpr(lbMap.getResult(0));
- AffineExpr ubExpr(ubMap.getResult(0));
- auto loopSpanExpr = simplifyAffineExpr(ubExpr - lbExpr, lbMap.getNumDims(),
- lbMap.getNumSymbols());
- auto cExpr = dyn_cast<AffineConstantExpr>(loopSpanExpr);
- if (!cExpr)
- return std::nullopt;
- return cExpr.getValue();
+ std::optional<uint64_t> tripCount;
+ for (AffineExpr ubExpr : ubMap.getResults()) {
+ AffineExpr loopSpanExpr = simplifyAffineExpr(
+ ubExpr - lbExpr, lbMap.getNumDims(), lbMap.getNumSymbols());
+ auto cExpr = dyn_cast<AffineConstantExpr>(loopSpanExpr);
+ if (!cExpr)
+ continue;
+ if (cExpr.getValue() < 0)
+ return 0;
+ tripCount =
+ std::min(tripCount.value_or(std::numeric_limits<uint64_t>::max()),
+ (uint64_t)cExpr.getValue());
+ }
+ return tripCount;
}
// Builds a map 'tripCountMap' from AffineForOp to constant trip count for loop
@@ -1899,7 +1914,8 @@ void mlir::affine::getComputationSliceState(
// Get bounds for slice IVs in terms of other IVs, symbols, and constants.
sliceCst.getSliceBounds(offset, numSliceLoopIVs, depSourceOp->getContext(),
- &sliceState->lbs, &sliceState->ubs);
+ &sliceState->lbs, &sliceState->ubs,
+ /*closedUb=*/false, /*allowMultiResultUb=*/true);
// Set up bound operands for the slice's lower and upper bounds.
SmallVector<Value, 4> sliceBoundOperands;
diff --git a/mlir/test/Dialect/Affine/loop-fusion-slice-computation.mlir b/mlir/test/Dialect/Affine/loop-fusion-slice-computation.mlir
index aa79ee26928e2..05bf3d415e1d6 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-slice-computation.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-slice-computation.mlir
@@ -160,3 +160,32 @@ func.func @slice_loop_nest_with_smaller_outer_trip_count() {
}
return
}
+
+// -----
+
+// The destination loop of a tiled dim is clamped at the end of the data by a
+// min upper bound, so the region of the source the tile reads is bounded by a
+// min too: `min(%i * 64 + 64, 1000)`. Such a bound is kept as the several
+// results it takes rather than being dropped for the constant bound, which
+// would say only that the slice ends somewhere before the end of the data and
+// make a slice of one tile look like a slice of everything from the tile on.
+
+#ub = affine_map<(d0) -> (64, d0 * -64 + 1000)>
+
+// CHECK-LABEL: func @slice_ub_from_min_bounded_dst() {
+func.func @slice_ub_from_min_bounded_dst() {
+ %0 = memref.alloc() : memref<1000xf32>
+ %cst = arith.constant 7.000000e+00 : f32
+ affine.for %i0 = 0 to 1000 {
+ // expected-remark at -1 {{Incorrect slice ( src loop: 1, dst loop: 0, depth: 1 : insert point: (1, 1) loop bounds: [(d0) -> (d0 floordiv 64), (d0) -> (d0 floordiv 64 + 1)] [(d0) -> (d0 mod 64), (d0) -> (d0 mod 64 + 1)] )}}
+ affine.store %cst, %0[%i0] : memref<1000xf32>
+ }
+ affine.for %i1 = 0 to 16 {
+ // expected-remark at -1 {{slice ( src loop: 0, dst loop: 1, depth: 1 : insert point: (1, 0) loop bounds: [(d0) -> (d0 * 64), (d0) -> (d0 * 64 + 64, 1000)] )}}
+ // expected-remark at -2 {{slice ( src loop: 0, dst loop: 1, depth: 2 : insert point: (2, 0) loop bounds: [(d0, d1) -> (d0 * 64 + d1), (d0, d1) -> (d0 * 64 + d1 + 1)] )}}
+ affine.for %i2 = 0 to min #ub(%i1) {
+ %1 = affine.load %0[%i1 * 64 + %i2] : memref<1000xf32>
+ }
+ }
+ return
+}
More information about the Mlir-commits
mailing list