[Mlir-commits] [mlir] [mlir][ValueBounds] Decouple merge collection and solving (PR #226944)
donald chen
llvmlistbot at llvm.org
Mon Sep 28 03:16:04 PDT 2026
https://github.com/cxy-1993 created https://github.com/llvm/llvm-project/pull/226944
Implement phase 1 of the ValueBounds constraint collection and solving RFC:
https://discourse.llvm.org/t/rfc-decoupling-constraint-collection-and-solving-in-mlir-valuebounds/91884
Introduce an atomic addMerge API that records the complete set of possible origins for a target. A target may register exactly one merge relationship. This prevents nested temporary queries from observing a partially collected relationship and moves conservative envelope computation out of individual op interfaces.
Expose ValueDim alongside the existing ValueDimList and migrate ValueBounds APIs to use it. A ValueDim is the graph-node identity used by ValueBounds, and representing both the target and every candidate this way allows each merge edge to select its own shaped dimension without positional Value/dimension argument pairs.
After processing the constraint worklist, solve collected merges to a fixed point. Single-candidate merges become equalities; multi-candidate merges derive independently provable lower and upper envelopes. Migrate arith.select and scf.if to relationship collection and add coverage for temporary evaluation, nested merges, single and multiple candidates, and differing candidate dimensions.
>From 0ae20001413f3a9a320a410a074287491de202df Mon Sep 17 00:00:00 2001
From: donald <chenxunyu1993 at gmail.com>
Date: Mon, 28 Sep 2026 10:10:05 +0000
Subject: [PATCH] [mlir][ValueBounds] Decouple merge collection and solving
Implement phase 1 of the ValueBounds constraint collection and solving RFC:
https://discourse.llvm.org/t/rfc-decoupling-constraint-collection-and-solving-in-mlir-valuebounds/91884
Introduce an atomic addMerge API that records the complete set of possible origins for a target. A target may register exactly one merge relationship. This prevents nested temporary queries from observing a partially collected relationship and moves conservative envelope computation out of individual op interfaces.
Expose ValueDim alongside the existing ValueDimList and migrate ValueBounds APIs to use it. A ValueDim is the graph-node identity used by ValueBounds, and representing both the target and every candidate this way allows each merge edge to select its own shaped dimension without positional Value/dimension argument pairs.
After processing the constraint worklist, solve collected merges to a fixed point. Single-candidate merges become equalities; multi-candidate merges derive independently provable lower and upper envelopes. Migrate arith.select and scf.if to relationship collection and add coverage for temporary evaluation, nested merges, single and multiple candidates, and differing candidate dimensions.
---
.../Dialect/Affine/Transforms/Transforms.h | 6 +-
.../IR/ScalableValueBoundsConstraintSet.h | 7 +-
.../mlir/Interfaces/ValueBoundsOpInterface.h | 107 +++++--
.../Affine/Transforms/ReifyValueBounds.cpp | 14 +-
.../Arith/IR/ValueBoundsOpInterfaceImpl.cpp | 52 +---
.../Arith/Transforms/ReifyValueBounds.cpp | 8 +-
.../IR/ValueBoundsOpInterfaceImpl.cpp | 2 +-
.../GPU/IR/ValueBoundsOpInterfaceImpl.cpp | 2 +-
.../Linalg/IR/ValueBoundsOpInterfaceImpl.cpp | 2 +-
.../Linalg/Transforms/HoistPadding.cpp | 3 +-
.../MemRef/IR/ValueBoundsOpInterfaceImpl.cpp | 12 +-
.../SCF/IR/ValueBoundsOpInterfaceImpl.cpp | 44 +--
.../Tensor/IR/ValueBoundsOpInterfaceImpl.cpp | 16 +-
.../IR/ScalableValueBoundsConstraintSet.cpp | 23 +-
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 3 +-
.../Transforms/VectorMaskElimination.cpp | 4 +-
.../lib/Interfaces/ValueBoundsOpInterface.cpp | 292 +++++++++++++-----
.../Arith/value-bounds-op-interface-impl.mlir | 75 +++++
.../SCF/value-bounds-op-interface-impl.mlir | 34 ++
.../Dialect/Affine/TestReifyValueBounds.cpp | 17 +-
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 31 ++
mlir/test/lib/Dialect/Test/TestOps.td | 26 ++
22 files changed, 527 insertions(+), 253 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h
index f8f6d4f5928b0..63f5217cd0ac8 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h
@@ -129,9 +129,9 @@ FailureOr<OpFoldResult> reifyShapedValueDimBound(
/// dialect ops.
/// * `reifyIndexValueBound`/`reifyShapedValueDimBound` are a combination of
/// the two functions mentioned above.
-OpFoldResult materializeComputedBound(
- OpBuilder &b, Location loc, AffineMap boundMap,
- ArrayRef<std::pair<Value, std::optional<int64_t>>> mapOperands);
+OpFoldResult materializeComputedBound(OpBuilder &b, Location loc,
+ AffineMap boundMap,
+ ArrayRef<ValueDim> mapOperands);
/// This transform tries to simplify the affine min operation `op`, by finding a
/// common lower bound for a set of expressions in the affine map results. It
diff --git a/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h b/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
index b32f4a777b90a..67391f3f70307 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
+++ b/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
@@ -35,7 +35,7 @@ struct ScalableValueBoundsConstraintSet
unsigned vscaleMin, unsigned vscaleMax, ValueBoundsOptions options = {})
: RTTIExtends(context, stopCondition, options,
/*addConservativeSemiAffineBounds=*/true),
- vscaleMin(vscaleMin), vscaleMax(vscaleMax) {};
+ vscaleMin(vscaleMin), vscaleMax(vscaleMax) {}
using RTTIExtends::bound;
using RTTIExtends::StopConditionFn;
@@ -69,9 +69,8 @@ struct ScalableValueBoundsConstraintSet
/// Note: `vscaleMin` must be `<=` to `vscaleMax`. If `vscaleMin` ==
/// `vscaleMax`, the resulting bound (if found), will be constant.
static FailureOr<ConstantOrScalableBound>
- computeScalableBound(Value value, std::optional<int64_t> dim,
- unsigned vscaleMin, unsigned vscaleMax,
- presburger::BoundType boundType,
+ computeScalableBound(ValueDim valueDim, unsigned vscaleMin,
+ unsigned vscaleMax, presburger::BoundType boundType,
ValueBoundsOptions options = {/*closedUB=*/true},
const StopConditionFn &stopCondition = nullptr);
diff --git a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
index 50c2dfe64de2b..0ee90a0eea374 100644
--- a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
+++ b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
@@ -48,9 +48,13 @@ class HyperrectangularSlice {
SmallVector<OpFoldResult> mixedStrides;
};
+/// An index/integer-typed SSA value (`std::nullopt`) or a dimension of a shaped
+/// SSA value. The dimension must be present if and only if the value is shaped.
+using ValueDim = std::pair<Value, std::optional<int64_t>>;
+
// Inline size chosen empirically based on compilation profiling.
// Profiled: 488K calls, avg=1.5+-0.5. N=2 covers >90% of cases inline.
-using ValueDimList = SmallVector<std::pair<Value, std::optional<int64_t>>, 2>;
+using ValueDimList = SmallVector<ValueDim, 2>;
/// Options that control value bound computation.
struct ValueBoundsOptions {
@@ -171,11 +175,10 @@ class ValueBoundsConstraintSet
/// index-type value. The traversal continues until the stop condition
/// evaluates to "true" for a value.
///
- /// The first parameter of the function is the shaped value/index-typed
- /// value. The second parameter is the dimension in case of a shaped value.
- /// The third parameter is this constraint set.
- using StopConditionFn = std::function<bool(
- Value, std::optional<int64_t> /*dim*/, ValueBoundsConstraintSet &cstr)>;
+ /// The first parameter is the shaped value dimension/index-typed value. The
+ /// second parameter is this constraint set.
+ using StopConditionFn =
+ std::function<bool(ValueDim, ValueBoundsConstraintSet &cstr)>;
/// Compute a bound for the given variable. The computed bound is stored in
/// `resultMap`. The operands of the bound are stored in `mapOperands`. An
@@ -239,20 +242,15 @@ class ValueBoundsConstraintSet
const StopConditionFn &stopCondition = nullptr,
ValueBoundsOptions options = {});
- /// Compute a constant delta between the given two values. Return "failure"
- /// if a constant delta could not be determined.
- ///
- /// `dim1`/`dim2` must be `nullopt` if and only if `value1`/`value2` are
- /// index-typed.
- static FailureOr<int64_t>
- computeConstantDelta(Value value1, Value value2,
- std::optional<int64_t> dim1 = std::nullopt,
- std::optional<int64_t> dim2 = std::nullopt);
+ /// Compute a constant delta between the given two values/dimensions. Return
+ /// "failure" if a constant delta could not be determined.
+ static FailureOr<int64_t> computeConstantDelta(ValueDim valueDim1,
+ ValueDim valueDim2);
/// Traverse the IR starting from the given value/dim and populate constraints
/// as long as the stop condition holds. Also process all values/dims that are
/// already on the worklist.
- void populateConstraints(Value value, std::optional<int64_t> dim);
+ void populateConstraints(ValueDim valueDim);
/// Comparison operator for `ValueBoundsConstraintSet::compare`.
enum ComparisonOperator { LT, LE, EQ, GT, GE };
@@ -340,8 +338,14 @@ class ValueBoundsConstraintSet
/// value dimension. If this value/dimension was not used so far, it is added
/// to the worklist.
///
- /// `dim` must be `nullopt` if and only if the given value is of index type.
- AffineExpr getExpr(Value value, std::optional<int64_t> dim = std::nullopt);
+ /// The dimension must be `nullopt` if and only if the given value is of index
+ /// type.
+ AffineExpr getExpr(ValueDim valueDim);
+
+ /// Convenience overload for an index-typed value.
+ AffineExpr getExpr(Value value) {
+ return getExpr(ValueDim{value, std::nullopt});
+ }
/// Return an expression that represents a constant or index-typed SSA value.
/// In case of a value, if this value was not used so far, it is added to the
@@ -351,18 +355,35 @@ class ValueBoundsConstraintSet
/// Return an expression that represents a constant.
AffineExpr getExpr(int64_t constant);
+ /// Record all possible origins of `value` as one complete relationship.
+ ///
+ /// A merge is a disjunction: `value` equals exactly one element of
+ /// `candidates` at runtime. The complete candidate list is passed atomically
+ /// so that temporary queries may solve the relationship immediately without
+ /// observing a partially constructed merge. An interface must call this
+ /// function exactly once for a given `value`.
+ ///
+ /// Candidate relationships are collected together with ordinary bounds and
+ /// solved after the current value-bounds worklist has been processed. This
+ /// keeps operation interface implementations declarative: an interface only
+ /// describes the possible origins, while the common merge solver decides
+ /// which conservative bounds can be derived from them.
+ ///
+ /// Each candidate carries its own dimension, independently of `valueDim`. A
+ /// merge may therefore relate different dimensions of shaped values, e.g.
+ /// result dimension 0 to candidate dimension 1. A dimension must be
+ /// `std::nullopt` if and only if its corresponding value is index-typed.
+ ///
+ /// A merge with exactly one candidate is an unconditional equality. Merges
+ /// with multiple candidates are solved conservatively after relationship
+ /// collection has finished.
+ void addMerge(ValueDim value, ValueDimList candidates);
+
/// Debugging only: Dump the constraint set and the column-to-value/dim
/// mapping to llvm::errs.
void dump() const;
protected:
- /// Dimension identifier to indicate a value is index-typed. This is used for
- /// internal data structures/API only.
- static constexpr int64_t kIndexValue = -1;
-
- /// An index-typed value or the dimension of a shaped-type value.
- using ValueDim = std::pair<Value, int64_t>;
-
ValueBoundsConstraintSet(MLIRContext *ctx,
const StopConditionFn &stopCondition,
ValueBoundsOptions options = {},
@@ -402,13 +423,21 @@ class ValueBoundsConstraintSet
/// any further.
void processWorklist();
+ /// Try to derive ordinary constraints from all collected candidate-origin
+ /// relationships. The phase 1 merge solver intentionally matches the
+ /// behavior of the merge logic that used to live in individual op interface
+ /// implementations. A single candidate becomes an equality. For multiple
+ /// candidates, any candidate proven to be below/above every other candidate
+ /// provides a conservative lower/upper envelope bound.
+ void solveMerges();
+
/// Bound the given column in the underlying constraint set by the given
/// expression.
void addBound(presburger::BoundType type, int64_t pos, AffineExpr expr);
/// Return the column position of the given value/dimension. Asserts that the
/// value/dimension exists in the constraint set.
- int64_t getPos(Value value, std::optional<int64_t> dim = std::nullopt) const;
+ int64_t getPos(ValueDim valueDim) const;
/// Return an affine expression that represents column `pos` in the constraint
/// set.
@@ -416,7 +445,7 @@ class ValueBoundsConstraintSet
/// Return "true" if the given value/dim is mapped (i.e., has a corresponding
/// column in the constraint system).
- bool isMapped(Value value, std::optional<int64_t> dim = std::nullopt) const;
+ bool isMapped(ValueDim valueDim) const;
/// Insert a value/dimension into the constraint set. If `isSymbol` is set to
/// "false", a dimension is added. The value/dimension is added to the
@@ -425,7 +454,7 @@ class ValueBoundsConstraintSet
/// Note: There are certain affine restrictions wrt. dimensions. E.g., they
/// cannot be multiplied. Furthermore, bounds can only be queried for
/// dimensions but not for symbols.
- int64_t insert(Value value, std::optional<int64_t> dim, bool isSymbol = true,
+ int64_t insert(ValueDim valueDim, bool isSymbol = true,
bool addToWorklist = true);
/// Insert an anonymous column into the constraint set. The column is not
@@ -465,6 +494,28 @@ class ValueBoundsConstraintSet
/// go stale.
std::queue<ValueDim> worklist;
+ /// A candidate-origin relationship. Unlike an ordinary equality, this
+ /// represents `value == candidate[0] OR value == candidate[1] OR ...`.
+ /// Keeping this relation separate from `cstr` is essential because
+ /// FlatLinearConstraints represents a conjunction of constraints and cannot
+ /// directly represent this disjunction.
+ struct MergeRelation {
+ ValueDim value;
+ ValueDimList candidates;
+
+ /// Lower and upper bounds are tracked separately because a partially
+ /// ordered candidate set may have a provable minimum but no provable
+ /// maximum, or vice versa. Once derived, a bound remains valid because the
+ /// constraint set only accumulates sound constraints.
+ bool lowerBoundSolved = false;
+ bool upperBoundSolved = false;
+ };
+
+ /// Candidate-origin relationships collected by op interfaces. Relationships
+ /// that cannot yet be solved are retained because solving another merge may
+ /// provide enough information to order their candidates on a later pass.
+ SmallVector<MergeRelation> mergeRelations;
+
/// Constraint system of equalities and inequalities.
FlatLinearConstraints cstr;
diff --git a/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp
index 1bd068d6151c6..b4717eee925a4 100644
--- a/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp
@@ -33,9 +33,9 @@ FailureOr<OpFoldResult> mlir::affine::reifyValueBound(
return affine::materializeComputedBound(b, loc, boundMap, mapOperands);
}
-OpFoldResult affine::materializeComputedBound(
- OpBuilder &b, Location loc, AffineMap boundMap,
- ArrayRef<std::pair<Value, std::optional<int64_t>>> mapOperands) {
+OpFoldResult affine::materializeComputedBound(OpBuilder &b, Location loc,
+ AffineMap boundMap,
+ ArrayRef<ValueDim> mapOperands) {
// Materialize tensor.dim/memref.dim ops.
SmallVector<Value> operands;
for (auto valueDim : mapOperands) {
@@ -89,14 +89,14 @@ FailureOr<OpFoldResult> mlir::affine::reifyShapedValueDimBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
int64_t dim, ValueBoundsConstraintSet::StopConditionFn stopCondition,
ValueBoundsOptions options) {
- auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
+ auto reifyToOperands = [&](ValueDim valueDim,
ValueBoundsConstraintSet &cstr) {
// We are trying to reify a bound for `value` in terms of the owning op's
// operands. Construct a stop condition that evaluates to "true" for any SSA
// value except for `value`. I.e., the bound will be computed in terms of
// any SSA values except for `value`. The first such values are operands of
// the owner of `value`.
- return v != value;
+ return valueDim.first != value;
};
return reifyValueBound(b, loc, type, {value, dim},
stopCondition ? stopCondition : reifyToOperands,
@@ -107,9 +107,9 @@ FailureOr<OpFoldResult> mlir::affine::reifyIndexValueBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
ValueBoundsConstraintSet::StopConditionFn stopCondition,
ValueBoundsOptions options) {
- auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
+ auto reifyToOperands = [&](ValueDim valueDim,
ValueBoundsConstraintSet &cstr) {
- return v != value;
+ return valueDim.first != value;
};
return reifyValueBound(b, loc, type, value,
stopCondition ? stopCondition : reifyToOperands,
diff --git a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
index 61e84e3cdae0f..958826a16cce3 100644
--- a/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -259,51 +259,19 @@ struct SelectOpInterface
if (isa<ShapedType>(condition.getType())) {
// If the condition is a shaped type, the condition is applied
// element-wise. All three operands must have the same shape.
- cstr.bound(value)[*dim] == cstr.getExpr(trueValue, dim);
- cstr.bound(value)[*dim] == cstr.getExpr(falseValue, dim);
- cstr.bound(value)[*dim] == cstr.getExpr(condition, dim);
+ cstr.bound(value)[*dim] == cstr.getExpr({trueValue, dim});
+ cstr.bound(value)[*dim] == cstr.getExpr({falseValue, dim});
+ cstr.bound(value)[*dim] == cstr.getExpr({condition, dim});
return;
}
- // Populate constraints for the true/false values (and all values on the
- // backward slice, as long as the current stop condition is not satisfied).
- cstr.populateConstraints(trueValue, dim);
- cstr.populateConstraints(falseValue, dim);
- auto boundsBuilder = cstr.bound(value);
- if (dim)
- boundsBuilder[*dim];
-
- // Compare yielded values.
- // If trueValue <= falseValue:
- // * result <= falseValue
- // * result >= trueValue
- if (cstr.populateAndCompare(
- /*lhs=*/{trueValue, dim},
- ValueBoundsConstraintSet::ComparisonOperator::LE,
- /*rhs=*/{falseValue, dim})) {
- if (dim) {
- cstr.bound(value)[*dim] >= cstr.getExpr(trueValue, dim);
- cstr.bound(value)[*dim] <= cstr.getExpr(falseValue, dim);
- } else {
- cstr.bound(value) >= trueValue;
- cstr.bound(value) <= falseValue;
- }
- }
- // If falseValue <= trueValue:
- // * result <= trueValue
- // * result >= falseValue
- if (cstr.populateAndCompare(
- /*lhs=*/{falseValue, dim},
- ValueBoundsConstraintSet::ComparisonOperator::LE,
- /*rhs=*/{trueValue, dim})) {
- if (dim) {
- cstr.bound(value)[*dim] >= cstr.getExpr(falseValue, dim);
- cstr.bound(value)[*dim] <= cstr.getExpr(trueValue, dim);
- } else {
- cstr.bound(value) >= falseValue;
- cstr.bound(value) <= trueValue;
- }
- }
+ // A scalar select has one of two candidate origins. Keep this disjunction
+ // explicit instead of comparing the candidates while the op interface is
+ // still collecting relationships. The common merge solver will reproduce
+ // the previous conservative envelope after candidate constraints have been
+ // collected.
+ ValueDimList candidates{{trueValue, dim}, {falseValue, dim}};
+ cstr.addMerge({value, dim}, std::move(candidates));
}
void populateBoundsForIndexValue(Operation *op, Value value,
diff --git a/mlir/lib/Dialect/Arith/Transforms/ReifyValueBounds.cpp b/mlir/lib/Dialect/Arith/Transforms/ReifyValueBounds.cpp
index 8a46762eb8e05..189b992943ff1 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ReifyValueBounds.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ReifyValueBounds.cpp
@@ -138,14 +138,14 @@ FailureOr<OpFoldResult> mlir::arith::reifyShapedValueDimBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
int64_t dim, const ValueBoundsConstraintSet::StopConditionFn &stopCondition,
ValueBoundsOptions options) {
- auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
+ auto reifyToOperands = [&](ValueDim valueDim,
ValueBoundsConstraintSet &cstr) {
// We are trying to reify a bound for `value` in terms of the owning op's
// operands. Construct a stop condition that evaluates to "true" for any SSA
// value expect for `value`. I.e., the bound will be computed in terms of
// any SSA values expect for `value`. The first such values are operands of
// the owner of `value`.
- return v != value;
+ return valueDim.first != value;
};
return reifyValueBound(b, loc, type, {value, dim},
stopCondition ? stopCondition : reifyToOperands,
@@ -156,9 +156,9 @@ FailureOr<OpFoldResult> mlir::arith::reifyIndexValueBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
const ValueBoundsConstraintSet::StopConditionFn &stopCondition,
ValueBoundsOptions options) {
- auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
+ auto reifyToOperands = [&](ValueDim valueDim,
ValueBoundsConstraintSet &cstr) {
- return v != value;
+ return valueDim.first != value;
};
return reifyValueBound(b, loc, type, value,
stopCondition ? stopCondition : reifyToOperands,
diff --git a/mlir/lib/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.cpp
index 8eec3fdec9f2f..bb2315659a4fe 100644
--- a/mlir/lib/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -35,7 +35,7 @@ struct MaterializationOpInterface
// shaped, for which no bound can be computed.
Value source = materializationOp.getOperand();
if (isa<ShapedType>(value.getType()) && isa<ShapedType>(source.getType()))
- cstr.bound(value)[dim] == cstr.getExpr(source, dim);
+ cstr.bound(value)[dim] == cstr.getExpr({source, dim});
}
};
diff --git a/mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp
index c978aefe3f655..67d5d55943844 100644
--- a/mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -58,7 +58,7 @@ struct SubgroupBroadcastOpInterface
ValueBoundsConstraintSet &cstr) const {
auto broadcastOp = cast<SubgroupBroadcastOp>(op);
assert(value == broadcastOp.getResult() && "invalid value");
- cstr.bound(value)[dim] == cstr.getExpr(broadcastOp.getSrc(), dim);
+ cstr.bound(value)[dim] == cstr.getExpr({broadcastOp.getSrc(), dim});
}
};
diff --git a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
index f56ef485069f8..40141d2a9c313 100644
--- a/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -39,7 +39,7 @@ struct IndexOpInterface
assert(flatDimPos >= flatDimCtr && "invalid pos");
auto shapedType = llvm::cast<ShapedType>(operand.getType());
if (flatDimPos < flatDimCtr + shapedType.getRank()) {
- cstr.bound(value) < cstr.getExpr(operand, flatDimPos - flatDimCtr);
+ cstr.bound(value) < cstr.getExpr({operand, flatDimPos - flatDimCtr});
break;
}
flatDimCtr += shapedType.getRank();
diff --git a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
index 85ec1c19384de..22f5dd7e48d40 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
@@ -473,7 +473,8 @@ HoistPaddingAnalysis::getHoistedPackedTensorSizes(RewriterBase &rewriter,
FailureOr<OpFoldResult> loopUb = affine::reifyIndexValueBound(
rewriter, loc, presburger::BoundType::UB, forOp.getUpperBound(),
/*stopCondition=*/
- [&](Value v, std::optional<int64_t> d, ValueBoundsConstraintSet &cstr) {
+ [&](ValueDim valueDim, ValueBoundsConstraintSet &cstr) {
+ Value v = valueDim.first;
if (v == forOp.getUpperBound())
return false;
// Compute a bound that is independent of any affine op results.
diff --git a/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
index 8408b166a6b2a..6187599a9e5b2 100644
--- a/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -39,7 +39,7 @@ struct AssumeAlignmentOpInterface
assert(value == assumeAlignmentOp.getResult() && "invalid value");
cstr.bound(value)[dim] ==
- cstr.getExpr(assumeAlignmentOp.getViewSource(), dim);
+ cstr.getExpr({assumeAlignmentOp.getViewSource(), dim});
}
};
@@ -52,7 +52,7 @@ struct CastOpInterface
if (llvm::isa<MemRefType>(castOp.getResult().getType()) &&
llvm::isa<MemRefType>(castOp.getSource().getType())) {
- cstr.bound(value)[dim] == cstr.getExpr(castOp.getSource(), dim);
+ cstr.bound(value)[dim] == cstr.getExpr({castOp.getSource(), dim});
}
}
};
@@ -68,7 +68,7 @@ struct DimOpInterface
auto constIndex = dimOp.getConstantIndex();
if (!constIndex.has_value())
return;
- cstr.bound(value) == cstr.getExpr(dimOp.getSource(), *constIndex);
+ cstr.bound(value) == cstr.getExpr({dimOp.getSource(), *constIndex});
}
};
@@ -104,7 +104,7 @@ struct ExtractStridedMetadataOpInterface
if (resultNumber >= sizeStart && resultNumber < strideStart) {
int64_t idx = resultNumber - sizeStart;
cstr.bound(value) >= 0;
- cstr.bound(value) == cstr.getExpr(metadataOp.getSource(), idx);
+ cstr.bound(value) == cstr.getExpr({metadataOp.getSource(), idx});
return;
}
@@ -160,10 +160,10 @@ struct CollapseShapeOpInterface
const ReassociationIndices reassocIndices =
collapseOp.getReassociationIndices()[dim];
AffineExpr productExpr =
- cstr.getExpr(collapseOp.getSrc(), reassocIndices[0]);
+ cstr.getExpr({collapseOp.getSrc(), reassocIndices[0]});
for (size_t i = 1; i < reassocIndices.size(); ++i) {
productExpr =
- productExpr * cstr.getExpr(collapseOp.getSrc(), reassocIndices[i]);
+ productExpr * cstr.getExpr({collapseOp.getSrc(), reassocIndices[i]});
}
cstr.bound(value)[dim] == productExpr;
}
diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
index 496a7b036e65d..b01ff7cb344b5 100644
--- a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -92,7 +92,7 @@ struct ForOpInterface
ValueBoundsConstraintSet::ComparisonOperator::EQ,
/*rhs=*/{iterArg, dim})) {
if (dim.has_value()) {
- cstr.bound(value)[*dim] == cstr.getExpr(initArg, dim);
+ cstr.bound(value)[*dim] == cstr.getExpr({initArg, dim});
} else {
cstr.bound(value) == cstr.getExpr(initArg);
}
@@ -168,7 +168,7 @@ struct ForallOpInterface
// The forall results and output arguments have the same sizes as the output
// operands.
Value outputOperand = forallOp.getOutputs()[iterArgIdx];
- cstr.bound(value)[dim] == cstr.getExpr(outputOperand, dim);
+ cstr.bound(value)[dim] == cstr.getExpr({outputOperand, dim});
}
};
@@ -182,41 +182,11 @@ struct IfOpInterface
Value thenValue = ifOp.thenYield().getResults()[resultNum];
Value elseValue = ifOp.elseYield().getResults()[resultNum];
- auto boundsBuilder = cstr.bound(value);
- if (dim)
- boundsBuilder[*dim];
-
- // Compare yielded values.
- // If thenValue <= elseValue:
- // * result <= elseValue
- // * result >= thenValue
- if (cstr.populateAndCompare(
- /*lhs=*/{thenValue, dim},
- ValueBoundsConstraintSet::ComparisonOperator::LE,
- /*rhs=*/{elseValue, dim})) {
- if (dim) {
- cstr.bound(value)[*dim] >= cstr.getExpr(thenValue, dim);
- cstr.bound(value)[*dim] <= cstr.getExpr(elseValue, dim);
- } else {
- cstr.bound(value) >= thenValue;
- cstr.bound(value) <= elseValue;
- }
- }
- // If elseValue <= thenValue:
- // * result <= thenValue
- // * result >= elseValue
- if (cstr.populateAndCompare(
- /*lhs=*/{elseValue, dim},
- ValueBoundsConstraintSet::ComparisonOperator::LE,
- /*rhs=*/{thenValue, dim})) {
- if (dim) {
- cstr.bound(value)[*dim] >= cstr.getExpr(elseValue, dim);
- cstr.bound(value)[*dim] <= cstr.getExpr(thenValue, dim);
- } else {
- cstr.bound(value) >= elseValue;
- cstr.bound(value) <= thenValue;
- }
- }
+ // Region results have one candidate origin per control-flow branch. Record
+ // that disjunction and let the common merge solver derive an envelope once
+ // both backward slices have been collected.
+ ValueDimList candidates{{thenValue, dim}, {elseValue, dim}};
+ cstr.addMerge({value, dim}, std::move(candidates));
}
void populateBoundsForIndexValue(Operation *op, Value value,
diff --git a/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp
index 684cec2d6748a..532ff48dd67d1 100644
--- a/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -27,7 +27,7 @@ struct CastOpInterface
if (llvm::isa<RankedTensorType>(castOp.getResult().getType()) &&
llvm::isa<RankedTensorType>(castOp.getSource().getType())) {
- cstr.bound(value)[dim] == cstr.getExpr(castOp.getSource(), dim);
+ cstr.bound(value)[dim] == cstr.getExpr({castOp.getSource(), dim});
}
}
};
@@ -44,10 +44,10 @@ struct CollapseShapeOpInterface
const ReassociationIndices reassocIndices =
collapseOp.getReassociationIndices()[dim];
AffineExpr productExpr =
- cstr.getExpr(collapseOp.getSrc(), reassocIndices[0]);
+ cstr.getExpr({collapseOp.getSrc(), reassocIndices[0]});
for (size_t i = 1; i < reassocIndices.size(); ++i) {
productExpr =
- productExpr * cstr.getExpr(collapseOp.getSrc(), reassocIndices[i]);
+ productExpr * cstr.getExpr({collapseOp.getSrc(), reassocIndices[i]});
}
cstr.bound(value)[dim] == productExpr;
}
@@ -67,14 +67,14 @@ struct ConcatOpInterface
// concatenated. Relate the result to every input: relating it to a single
// input loses the bound when that input is the unbounded one.
for (Value input : inputs)
- cstr.bound(value)[dim] == cstr.getExpr(input, dim);
+ cstr.bound(value)[dim] == cstr.getExpr({input, dim});
return;
}
// The concatenated dimension is the sum of the input sizes.
- AffineExpr sum = cstr.getExpr(inputs.front(), dim);
+ AffineExpr sum = cstr.getExpr({inputs.front(), dim});
for (Value input : inputs.drop_front())
- sum = sum + cstr.getExpr(input, dim);
+ sum = sum + cstr.getExpr({input, dim});
cstr.bound(value)[dim] == sum;
}
};
@@ -90,7 +90,7 @@ struct DimOpInterface
auto constIndex = dimOp.getConstantIndex();
if (!constIndex.has_value())
return;
- cstr.bound(value) == cstr.getExpr(dimOp.getSource(), *constIndex);
+ cstr.bound(value) == cstr.getExpr({dimOp.getSource(), *constIndex});
}
};
@@ -146,7 +146,7 @@ struct PadOpInterface
auto padOp = cast<PadOp>(op);
assert(value == padOp.getResult() && "invalid value");
- AffineExpr srcSize = cstr.getExpr(padOp.getSource(), dim);
+ AffineExpr srcSize = cstr.getExpr({padOp.getSource(), dim});
AffineExpr lowPad = cstr.getExpr(padOp.getMixedLowPad()[dim]);
AffineExpr highPad = cstr.getExpr(padOp.getMixedHighPad()[dim]);
cstr.bound(value)[dim] == srcSize + lowPad + highPad;
diff --git a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
index ed41d19d627a3..934caadee2a68 100644
--- a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
+++ b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
@@ -41,15 +41,17 @@ char ScalableValueBoundsConstraintSet::ID = 0;
FailureOr<ConstantOrScalableBound>
ScalableValueBoundsConstraintSet::computeScalableBound(
- Value value, std::optional<int64_t> dim, unsigned vscaleMin,
- unsigned vscaleMax, presburger::BoundType boundType,
- ValueBoundsOptions options, const StopConditionFn &stopCondition) {
+ ValueDim valueDim, unsigned vscaleMin, unsigned vscaleMax,
+ presburger::BoundType boundType, ValueBoundsOptions options,
+ const StopConditionFn &stopCondition) {
using namespace presburger;
assert(vscaleMin <= vscaleMax);
+ Value value = valueDim.first;
+ std::optional<int64_t> dim = valueDim.second;
// No stop condition specified: Keep adding constraints until the worklist
// is empty.
- auto defaultStopCondition = [&](Value v, std::optional<int64_t> dim,
+ auto defaultStopCondition = [&](ValueDim valueDim,
mlir::ValueBoundsConstraintSet &cstr) {
return false;
};
@@ -57,7 +59,7 @@ ScalableValueBoundsConstraintSet::computeScalableBound(
ScalableValueBoundsConstraintSet scalableCstr(
value.getContext(), stopCondition ? stopCondition : defaultStopCondition,
vscaleMin, vscaleMax, options);
- int64_t pos = scalableCstr.insert(value, dim, /*isSymbol=*/false);
+ int64_t pos = scalableCstr.insert(valueDim, /*isSymbol=*/false);
scalableCstr.processWorklist();
// Check the resulting constraints set is valid.
@@ -67,11 +69,9 @@ ScalableValueBoundsConstraintSet::computeScalableBound(
// Project out all columns apart from vscale and the starting point
// (value/dim). This should result in constraints in terms of vscale only.
- auto projectOutFn = [&](ValueDim p) {
- bool isStartingPoint =
- p.first == value &&
- p.second == dim.value_or(ValueBoundsConstraintSet::kIndexValue);
- return p.first != scalableCstr.getVscaleValue() && !isStartingPoint;
+ auto projectOutFn = [&](ValueDim candidate) {
+ bool isStartingPoint = candidate == valueDim;
+ return candidate.first != scalableCstr.getVscaleValue() && !isStartingPoint;
};
scalableCstr.projectOut(projectOutFn);
scalableCstr.projectOutAnonymous(/*except=*/pos);
@@ -90,8 +90,7 @@ ScalableValueBoundsConstraintSet::computeScalableBound(
if (i == pos)
continue;
if (scalableCstr.positionToValueDim[i] !=
- ValueDim(scalableCstr.getVscaleValue(),
- ValueBoundsConstraintSet::kIndexValue)) {
+ ValueDim(scalableCstr.getVscaleValue(), std::nullopt)) {
return failure();
}
}
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index d5fd5095b18b0..42367ca410334 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -308,7 +308,8 @@ bool mlir::vector::isDisjointTransferIndices(
return true;
FailureOr<int64_t> computeDelta =
- ValueBoundsConstraintSet::computeConstantDelta(indexA, indexB);
+ ValueBoundsConstraintSet::computeConstantDelta(
+ {indexA, std::nullopt}, {indexB, std::nullopt});
if (succeeded(computeDelta)) {
if (std::abs(computeDelta.value()) >= vectorDim)
return true;
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorMaskElimination.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorMaskElimination.cpp
index 6f75ce7a04511..c118b94311237 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorMaskElimination.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorMaskElimination.cpp
@@ -56,8 +56,8 @@ LogicalResult resolveAllTrueCreateMaskOp(IRRewriter &rewriter,
// value it could be).
FailureOr<ConstantOrScalableBound> dimLowerBound =
vector::ScalableValueBoundsConstraintSet::computeScalableBound(
- dimSize, {}, vscaleRange.vscaleMin, vscaleRange.vscaleMax,
- presburger::BoundType::LB);
+ {dimSize, std::nullopt}, vscaleRange.vscaleMin,
+ vscaleRange.vscaleMax, presburger::BoundType::LB);
if (failed(dimLowerBound))
return failure();
auto dimLowerBoundSize = dimLowerBound->getSize();
diff --git a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
index d6d3b4d43cb55..e3069619524ba 100644
--- a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
+++ b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
@@ -205,8 +205,8 @@ ValueBoundsConstraintSet::ValueBoundsConstraintSet(
char ValueBoundsConstraintSet::ID = 0;
#ifndef NDEBUG
-static void assertValidValueDim(Value value, std::optional<int64_t> dim,
- ValueBoundsOptions options) {
+static void assertValidValueDim(ValueDim valueDim, ValueBoundsOptions options) {
+ auto [value, dim] = valueDim;
if (isIndexLikeType(value.getType(), options)) {
assert(!dim.has_value() && "invalid dim value");
} else if (auto shapedType = dyn_cast<ShapedType>(value.getType())) {
@@ -242,10 +242,10 @@ void ValueBoundsConstraintSet::addBound(BoundType type, int64_t pos,
}
}
-AffineExpr ValueBoundsConstraintSet::getExpr(Value value,
- std::optional<int64_t> dim) {
+AffineExpr ValueBoundsConstraintSet::getExpr(ValueDim valueDim) {
+ auto [value, dim] = valueDim;
#ifndef NDEBUG
- assertValidValueDim(value, dim, options);
+ assertValidValueDim(valueDim, options);
#endif // NDEBUG
// Check if the value/dim is statically known. In that case, an affine
@@ -263,21 +263,20 @@ AffineExpr ValueBoundsConstraintSet::getExpr(Value value,
// If the value/dim is already mapped, return the corresponding expression
// directly.
- ValueDim valueDim = std::make_pair(value, dim.value_or(kIndexValue));
if (valueDimToPosition.contains(valueDim)) {
// If it is a constant, return an affine constant expression. Otherwise,
// return an affine expression that represents the respective column in the
// constraint set.
if (constSize)
return builder.getAffineConstantExpr(*constSize);
- return getPosExpr(getPos(value, dim));
+ return getPosExpr(getPos(valueDim));
}
if (constSize) {
// Constant index value/dim: add column to the constraint set, add EQ bound
// and return an affine constant expression without pushing the newly added
// column to the worklist.
- (void)insert(value, dim, /*isSymbol=*/true, /*addToWorklist=*/false);
+ (void)insert(valueDim, /*isSymbol=*/true, /*addToWorklist=*/false);
if (shapedType)
bound(value)[*dim] == *constSize;
else
@@ -288,12 +287,12 @@ AffineExpr ValueBoundsConstraintSet::getExpr(Value value,
// Dynamic value/dim: insert column to the constraint set and put it on the
// worklist. Return an affine expression that represents the newly inserted
// column in the constraint set.
- return getPosExpr(insert(value, dim, /*isSymbol=*/true));
+ return getPosExpr(insert(valueDim, /*isSymbol=*/true));
}
AffineExpr ValueBoundsConstraintSet::getExpr(OpFoldResult ofr) {
if (Value value = llvm::dyn_cast_if_present<Value>(ofr))
- return getExpr(value, /*dim=*/std::nullopt);
+ return getExpr(value);
auto constInt = ::getConstantIntValue(ofr);
assert(constInt.has_value() && "expected Integer constant");
return builder.getAffineConstantExpr(*constInt);
@@ -303,20 +302,54 @@ AffineExpr ValueBoundsConstraintSet::getExpr(int64_t constant) {
return builder.getAffineConstantExpr(constant);
}
-int64_t ValueBoundsConstraintSet::insert(Value value,
- std::optional<int64_t> dim,
- bool isSymbol, bool addToWorklist) {
+void ValueBoundsConstraintSet::addMerge(ValueDim valueDim,
+ ValueDimList candidates) {
+#ifndef NDEBUG
+ assertValidValueDim(valueDim, options);
+#endif // NDEBUG
+ assert(isMapped(valueDim) &&
+ "the merged value must already be in the constraint set");
+ assert(!candidates.empty() && "expected at least one merge candidate");
+ assert(llvm::none_of(mergeRelations,
+ [&](const MergeRelation &relation) {
+ return relation.value == valueDim;
+ }) &&
+ "a merge relationship was already registered for this value");
+
+ ValueDimList candidateValues;
+ candidateValues.reserve(candidates.size());
+ for (auto [candidate, candidateDim] : candidates) {
+#ifndef NDEBUG
+ assertValidValueDim({candidate, candidateDim}, options);
+#endif // NDEBUG
+
+ // Map every candidate and schedule its defining relationship for
+ // collection. No merge-derived constraints are added until the complete
+ // relation has been stored below.
+ (void)getExpr({candidate, candidateDim});
+ ValueDim candidateValue{candidate, candidateDim};
+ if (!llvm::is_contained(candidateValues, candidateValue))
+ candidateValues.push_back(candidateValue);
+ }
+
+ mergeRelations.push_back(
+ MergeRelation{/*value=*/valueDim,
+ /*candidates=*/std::move(candidateValues)});
+}
+
+int64_t ValueBoundsConstraintSet::insert(ValueDim valueDim, bool isSymbol,
+ bool addToWorklist) {
+ auto [value, dim] = valueDim;
#ifndef NDEBUG
- assertValidValueDim(value, dim, options);
+ assertValidValueDim(valueDim, options);
#endif // NDEBUG
- ValueDim valueDim = std::make_pair(value, dim.value_or(kIndexValue));
assert(!valueDimToPosition.contains(valueDim) && "already mapped");
int64_t pos = isSymbol ? cstr.appendVar(VarKind::Symbol)
: cstr.appendVar(VarKind::SetDim);
LDBG() << "Inserting constraint set column " << pos << " for: " << value
- << " (dim: " << dim.value_or(kIndexValue)
- << ", owner: " << getOwnerOfValue(value)->getName() << ")";
+ << " (dim: " << dim << ", owner: " << getOwnerOfValue(value)->getName()
+ << ")";
positionToValueDim.insert(positionToValueDim.begin() + pos, valueDim);
// Update reverse mapping.
for (int64_t i = pos, e = positionToValueDim.size(); i < e; ++i)
@@ -333,8 +366,7 @@ int64_t ValueBoundsConstraintSet::insert(Value value,
if (addToWorklist &&
(!isa<BlockArgument>(value) ||
cast<BlockArgument>(value).getOwner()->isEntryBlock())) {
- LDBG() << "Push to worklist: " << value
- << " (dim: " << dim.value_or(kIndexValue) << ")";
+ LDBG() << "Push to worklist: " << value << " (dim: " << dim << ")";
worklist.push(valueDim);
}
@@ -362,9 +394,7 @@ int64_t ValueBoundsConstraintSet::insert(AffineMap map,
// Add map and operands to the constraint set. Dimensions are converted to
// symbols. All operands are added to the worklist (unless they were already
// processed).
- auto mapper = [&](std::pair<Value, std::optional<int64_t>> v) {
- return getExpr(v.first, v.second);
- };
+ auto mapper = [&](ValueDim valueDim) { return getExpr(valueDim); };
SmallVector<AffineExpr> dimReplacements = llvm::map_to_vector(
ArrayRef(operands).take_front(map.getNumDims()), mapper);
SmallVector<AffineExpr> symReplacements = llvm::map_to_vector(
@@ -380,16 +410,14 @@ int64_t ValueBoundsConstraintSet::insert(const Variable &var, bool isSymbol) {
return insert(var.map, var.mapOperands, isSymbol);
}
-int64_t ValueBoundsConstraintSet::getPos(Value value,
- std::optional<int64_t> dim) const {
+int64_t ValueBoundsConstraintSet::getPos(ValueDim valueDim) const {
+ auto [value, dim] = valueDim;
#ifndef NDEBUG
- assertValidValueDim(value, dim, options);
+ assertValidValueDim(valueDim, options);
#endif // NDEBUG
- LDBG() << "Getting pos for: " << value
- << " (dim: " << dim.value_or(kIndexValue)
+ LDBG() << "Getting pos for: " << value << " (dim: " << dim
<< ", owner: " << getOwnerOfValue(value)->getName() << ")";
- auto it =
- valueDimToPosition.find(std::make_pair(value, dim.value_or(kIndexValue)));
+ auto it = valueDimToPosition.find(valueDim);
assert(it != valueDimToPosition.end() && "expected mapped entry");
return it->second;
}
@@ -401,11 +429,8 @@ AffineExpr ValueBoundsConstraintSet::getPosExpr(int64_t pos) {
: builder.getAffineSymbolExpr(pos - cstr.getNumDimVars());
}
-bool ValueBoundsConstraintSet::isMapped(Value value,
- std::optional<int64_t> dim) const {
- auto it =
- valueDimToPosition.find(std::make_pair(value, dim.value_or(kIndexValue)));
- return it != valueDimToPosition.end();
+bool ValueBoundsConstraintSet::isMapped(ValueDim valueDim) const {
+ return valueDimToPosition.contains(valueDim);
}
void ValueBoundsConstraintSet::processWorklist() {
@@ -416,22 +441,20 @@ void ValueBoundsConstraintSet::processWorklist() {
assert(valueDimToPosition.contains(valueDim) &&
"expected mapped worklist entry");
Value value = valueDim.first;
- int64_t dim = valueDim.second;
+ std::optional<int64_t> dim = valueDim.second;
// Check for static dim size.
- if (dim != kIndexValue) {
+ if (dim) {
auto shapedType = cast<ShapedType>(value.getType());
- if (shapedType.hasRank() && !shapedType.isDynamicDim(dim)) {
- bound(value)[dim] == getExpr(shapedType.getDimSize(dim));
+ if (shapedType.hasRank() && !shapedType.isDynamicDim(*dim)) {
+ bound(value)[*dim] == getExpr(shapedType.getDimSize(*dim));
continue;
}
}
// Do not process any further if the stop condition is met.
- auto maybeDim = dim == kIndexValue ? std::nullopt : std::make_optional(dim);
- if (stopCondition(value, maybeDim, *this)) {
- LDBG() << "Stop condition met for: " << value << " (dim: " << maybeDim
- << ")";
+ if (stopCondition(valueDim, *this)) {
+ LDBG() << "Stop condition met for: " << value << " (dim: " << dim << ")";
continue;
}
@@ -442,10 +465,10 @@ void ValueBoundsConstraintSet::processWorklist() {
LDBG() << "Query value bounds for: " << value
<< " (owner: " << getOwnerOfValue(value)->getName() << ")";
if (valueBoundsOp) {
- if (dim == kIndexValue) {
+ if (!dim) {
valueBoundsOp.populateBoundsForIndexValue(value, *this);
} else {
- valueBoundsOp.populateBoundsForShapedValueDim(value, dim, *this);
+ valueBoundsOp.populateBoundsForShapedValueDim(value, *dim, *this);
}
continue;
}
@@ -455,11 +478,114 @@ void ValueBoundsConstraintSet::processWorklist() {
// implements the `DestinationStyleOpInterface`. OpResults of such ops are
// tied to OpOperands. Tied values have the same shape.
auto dstOp = value.getDefiningOp<DestinationStyleOpInterface>();
- if (!dstOp || dim == kIndexValue)
+ if (!dstOp || !dim)
continue;
Value tiedOperand = dstOp.getTiedOpOperand(cast<OpResult>(value))->get();
- bound(value)[dim] == getExpr(tiedOperand, dim);
+ bound(value)[*dim] == getExpr({tiedOperand, dim});
}
+
+ // Every merge is registered atomically with its complete candidate list, so
+ // it is safe to solve here even when processWorklist was entered recursively
+ // by a temporary query such as populateAndCompare. This is required for that
+ // query to observe bounds derived from relationships it just collected.
+ solveMerges();
+}
+
+void ValueBoundsConstraintSet::solveMerges() {
+ // Solving one merge can make another merge solvable. For example, in
+ //
+ // inner = merge(4, 9)
+ // outer = merge(inner, 12)
+ //
+ // the bounds of `inner` must be added before `inner` can be compared with
+ // 12. Relationships are discovered in worklist order, which is not a
+ // topological order in the presence of shared dependencies. Iterate over the
+ // unsolved relations until a pass makes no progress. Every successful
+ // relation is marked solved, so this loop performs at most N successful
+ // iterations for the two independently tracked bounds of every relation.
+ bool madeProgress;
+ do {
+ madeProgress = false;
+ for (MergeRelation &relation : mergeRelations) {
+ if (relation.lowerBoundSolved && relation.upperBoundSolved)
+ continue;
+
+ // No interface should normally create an empty relation because a
+ // relation is allocated by addMerge together with its first candidate.
+ // Keep the guard to make this routine robust against future construction
+ // paths that may reserve a merge before discovering its predecessors.
+ if (relation.candidates.empty())
+ continue;
+
+ auto getValueExpr = [&](ValueDim valueDim) { return getExpr(valueDim); };
+ auto addResultBound = [&](presburger::BoundType type,
+ ValueDim candidate) {
+ std::optional<int64_t> resultDim = relation.value.second;
+ BoundBuilder resultBound = bound(relation.value.first);
+ if (resultDim)
+ resultBound[*resultDim];
+ if (type == presburger::BoundType::LB)
+ resultBound >= getValueExpr(candidate);
+ else
+ resultBound <= getValueExpr(candidate);
+ };
+
+ // A one-candidate merge is no longer a disjunction. Preserve that exact
+ // information as an equality instead of weakening it to an envelope.
+ if (relation.candidates.size() == 1) {
+ ValueDim candidate = relation.candidates.front();
+ std::optional<int64_t> resultDim = relation.value.second;
+ BoundBuilder resultBound = bound(relation.value.first);
+ if (resultDim)
+ resultBound[*resultDim];
+ resultBound == getValueExpr(candidate);
+ relation.lowerBoundSolved = true;
+ relation.upperBoundSolved = true;
+ madeProgress = true;
+ continue;
+ }
+
+ // Find a candidate that is no greater than every other candidate. Such a
+ // candidate is a sound lower bound for the merge, regardless of which
+ // origin is selected at runtime. This O(N^2) search intentionally does
+ // not require a total order: it can still produce one side of the
+ // envelope for a partially ordered set.
+ auto findExtremum =
+ [&](ComparisonOperator comparison) -> std::optional<ValueDim> {
+ for (ValueDim candidate : relation.candidates) {
+ int64_t candidatePos = getPos(candidate);
+ bool isExtremum =
+ llvm::all_of(relation.candidates, [&](ValueDim otherCandidate) {
+ if (candidate == otherCandidate)
+ return true;
+ int64_t otherPos = getPos(otherCandidate);
+ return comparePos(candidatePos, comparison, otherPos);
+ });
+ if (isExtremum)
+ return candidate;
+ }
+ return std::nullopt;
+ };
+
+ if (!relation.lowerBoundSolved) {
+ if (std::optional<ValueDim> lower =
+ findExtremum(ComparisonOperator::LE)) {
+ addResultBound(presburger::BoundType::LB, *lower);
+ relation.lowerBoundSolved = true;
+ madeProgress = true;
+ }
+ }
+
+ if (!relation.upperBoundSolved) {
+ if (std::optional<ValueDim> upper =
+ findExtremum(ComparisonOperator::GE)) {
+ addResultBound(presburger::BoundType::UB, *upper);
+ relation.upperBoundSolved = true;
+ madeProgress = true;
+ }
+ }
+ }
+ } while (madeProgress);
}
void ValueBoundsConstraintSet::projectOut(int64_t pos) {
@@ -525,11 +651,8 @@ LogicalResult ValueBoundsConstraintSet::computeBound(
// Project out all variables (apart from `valueDim`) that do not match the
// stop condition.
- cstr.projectOut([&](ValueDim p) {
- auto maybeDim =
- p.second == kIndexValue ? std::nullopt : std::make_optional(p.second);
- return !stopCondition(p.first, maybeDim, cstr);
- });
+ cstr.projectOut(
+ [&](ValueDim valueDim) { return !stopCondition(valueDim, cstr); });
cstr.projectOutAnonymous(/*except=*/pos);
// Compute lower and upper bounds for `valueDim`.
@@ -608,10 +731,10 @@ LogicalResult ValueBoundsConstraintSet::computeBound(
assert(cstr.positionToValueDim[i].has_value() &&
"cannot build affine map in terms of anonymous column");
- ValueBoundsConstraintSet::ValueDim valueDim = *cstr.positionToValueDim[i];
+ ValueDim valueDim = *cstr.positionToValueDim[i];
Value value = valueDim.first;
- int64_t dim = valueDim.second;
- if (dim == ValueBoundsConstraintSet::kIndexValue) {
+ std::optional<int64_t> dim = valueDim.second;
+ if (!dim) {
// An index-typed/integer-typed value is used: it can be used directly in
// the computed bound.
assert(isIndexLikeType(value.getType(), options) &&
@@ -620,9 +743,9 @@ LogicalResult ValueBoundsConstraintSet::computeBound(
continue;
}
- assert(cast<ShapedType>(value.getType()).isDynamicDim(dim) &&
+ assert(cast<ShapedType>(value.getType()).isDynamicDim(*dim) &&
"expected dynamic dim");
- mapOperands.push_back(std::make_pair(value, dim));
+ mapOperands.push_back(valueDim);
}
resultMap = bound.replaceDimsAndSymbols(replacementDims, replacementSymbols,
@@ -636,8 +759,8 @@ LogicalResult ValueBoundsConstraintSet::computeDependentBound(
ValueBoundsOptions options) {
return computeBound(
resultMap, mapOperands, type, var,
- [&](Value v, std::optional<int64_t> d, ValueBoundsConstraintSet &cstr) {
- return llvm::is_contained(dependencies, std::make_pair(v, d));
+ [&](ValueDim valueDim, ValueBoundsConstraintSet &cstr) {
+ return llvm::is_contained(dependencies, valueDim);
},
options);
}
@@ -671,8 +794,8 @@ LogicalResult ValueBoundsConstraintSet::computeIndependentBound(
// Reify bounds in terms of any independent values.
return computeBound(
resultMap, mapOperands, type, var,
- [&](Value v, std::optional<int64_t> d, ValueBoundsConstraintSet &cstr) {
- return isIndependent(v);
+ [&](ValueDim valueDim, ValueBoundsConstraintSet &cstr) {
+ return isIndependent(valueDim.first);
},
options);
}
@@ -683,7 +806,7 @@ FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
// Default stop condition if none was specified: Keep adding constraints until
// a bound could be computed.
int64_t pos = 0;
- auto defaultStopCondition = [&](Value v, std::optional<int64_t> dim,
+ auto defaultStopCondition = [&](ValueDim valueDim,
ValueBoundsConstraintSet &cstr) {
return cstr.cstr.getConstantBound64(type, pos).has_value();
};
@@ -701,15 +824,14 @@ FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
return failure();
}
-void ValueBoundsConstraintSet::populateConstraints(Value value,
- std::optional<int64_t> dim) {
+void ValueBoundsConstraintSet::populateConstraints(ValueDim valueDim) {
#ifndef NDEBUG
- assertValidValueDim(value, dim, options);
+ assertValidValueDim(valueDim, options);
#endif // NDEBUG
// `getExpr` pushes the value/dim onto the worklist (unless it was already
// analyzed).
- (void)getExpr(value, dim);
+ (void)getExpr(valueDim);
// Process all values/dims on the worklist. This may traverse and analyze
// additional IR, depending the current stop function.
processWorklist();
@@ -725,19 +847,21 @@ int64_t ValueBoundsConstraintSet::populateConstraints(AffineMap map,
}
FailureOr<int64_t>
-ValueBoundsConstraintSet::computeConstantDelta(Value value1, Value value2,
- std::optional<int64_t> dim1,
- std::optional<int64_t> dim2) {
+ValueBoundsConstraintSet::computeConstantDelta(ValueDim valueDim1,
+ ValueDim valueDim2) {
#ifndef NDEBUG
- assertValidValueDim(value1, dim1, /*options=*/{});
- assertValidValueDim(value2, dim2, /*options=*/{});
+ assertValidValueDim(valueDim1, /*options=*/{});
+ assertValidValueDim(valueDim2, /*options=*/{});
#endif // NDEBUG
- Builder b(value1.getContext());
+ Builder b(valueDim1.first.getContext());
AffineMap map = AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0,
b.getAffineDimExpr(0) - b.getAffineDimExpr(1));
+ SmallVector<Variable, 2> operands;
+ operands.emplace_back(valueDim1.first, valueDim1.second);
+ operands.emplace_back(valueDim2.first, valueDim2.second);
return computeConstantBound(presburger::BoundType::EQ,
- Variable(map, {{value1, dim1}, {value2, dim2}}));
+ Variable(map, operands));
}
bool ValueBoundsConstraintSet::comparePos(int64_t lhsPos,
@@ -837,8 +961,7 @@ bool ValueBoundsConstraintSet::compare(const Variable &lhs,
ComparisonOperator cmp,
const Variable &rhs) {
int64_t lhsPos = -1, rhsPos = -1;
- auto stopCondition = [&](Value v, std::optional<int64_t> dim,
- ValueBoundsConstraintSet &cstr) {
+ auto stopCondition = [&](ValueDim valueDim, ValueBoundsConstraintSet &cstr) {
// Keep processing as long as lhs/rhs were not processed.
if (size_t(lhsPos) >= cstr.positionToValueDim.size() ||
size_t(rhsPos) >= cstr.positionToValueDim.size())
@@ -854,8 +977,7 @@ FailureOr<bool> ValueBoundsConstraintSet::strongCompare(const Variable &lhs,
ComparisonOperator cmp,
const Variable &rhs) {
int64_t lhsPos = -1, rhsPos = -1;
- auto stopCondition = [&](Value v, std::optional<int64_t> dim,
- ValueBoundsConstraintSet &cstr) {
+ auto stopCondition = [&](ValueDim valueDim, ValueBoundsConstraintSet &cstr) {
// Keep processing as long as lhs/rhs were not processed.
if (size_t(lhsPos) >= cstr.positionToValueDim.size() ||
size_t(rhsPos) >= cstr.positionToValueDim.size())
@@ -986,10 +1108,10 @@ void ValueBoundsConstraintSet::dump() const {
for (auto [index, valueDim] : llvm::enumerate(positionToValueDim)) {
llvm::errs() << " " << index << "\t";
if (valueDim) {
- if (valueDim->second == kIndexValue) {
+ if (!valueDim->second) {
llvm::errs() << "n/a\t";
} else {
- llvm::errs() << valueDim->second << "\t";
+ llvm::errs() << *valueDim->second << "\t";
}
llvm::errs() << getOwnerOfValue(valueDim->first)->getName() << " ";
if (OpResult result = dyn_cast<OpResult>(valueDim->first)) {
@@ -1014,16 +1136,16 @@ ValueBoundsConstraintSet::BoundBuilder::operator[](int64_t dim) {
assert(!this->dim.has_value() && "dim was already set");
this->dim = dim;
#ifndef NDEBUG
- assertValidValueDim(value, this->dim, cstr.options);
+ assertValidValueDim({value, this->dim}, cstr.options);
#endif // NDEBUG
return *this;
}
void ValueBoundsConstraintSet::BoundBuilder::operator<(AffineExpr expr) {
#ifndef NDEBUG
- assertValidValueDim(value, this->dim, cstr.options);
+ assertValidValueDim({value, this->dim}, cstr.options);
#endif // NDEBUG
- cstr.addBound(BoundType::UB, cstr.getPos(value, this->dim), expr);
+ cstr.addBound(BoundType::UB, cstr.getPos({value, this->dim}), expr);
}
void ValueBoundsConstraintSet::BoundBuilder::operator<=(AffineExpr expr) {
@@ -1036,16 +1158,16 @@ void ValueBoundsConstraintSet::BoundBuilder::operator>(AffineExpr expr) {
void ValueBoundsConstraintSet::BoundBuilder::operator>=(AffineExpr expr) {
#ifndef NDEBUG
- assertValidValueDim(value, this->dim, cstr.options);
+ assertValidValueDim({value, this->dim}, cstr.options);
#endif // NDEBUG
- cstr.addBound(BoundType::LB, cstr.getPos(value, this->dim), expr);
+ cstr.addBound(BoundType::LB, cstr.getPos({value, this->dim}), expr);
}
void ValueBoundsConstraintSet::BoundBuilder::operator==(AffineExpr expr) {
#ifndef NDEBUG
- assertValidValueDim(value, this->dim, cstr.options);
+ assertValidValueDim({value, this->dim}, cstr.options);
#endif // NDEBUG
- cstr.addBound(BoundType::EQ, cstr.getPos(value, this->dim), expr);
+ cstr.addBound(BoundType::EQ, cstr.getPos({value, this->dim}), expr);
}
void ValueBoundsConstraintSet::BoundBuilder::operator<(OpFoldResult ofr) {
diff --git a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
index ceb8256003b1e..45bf2d9c4abab 100644
--- a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
@@ -165,6 +165,25 @@ func.func @arith_divui_positive_lhs() -> index {
}
+// -----
+
+// The divui interface temporarily queries whether its lhs is non-negative.
+// That nested query must solve the complete select merge before returning;
+// otherwise divui would add no relationship for its result.
+// CHECK-LABEL: func @arith_divui_merge_lhs(
+// CHECK: %[[c3:.*]] = arith.constant 3 : index
+// CHECK: return %[[c3]]
+func.func @arith_divui_merge_lhs(%cond: i1) -> index {
+ %c6 = arith.constant 6 : index
+ %c10 = arith.constant 10 : index
+ %c2 = arith.constant 2 : index
+ %lhs = arith.select %cond, %c6, %c10 : index
+ %quotient = arith.divui %lhs, %c2 : index
+ %lb = "test.reify_bound"(%quotient) {type = "LB", constant}
+ : (index) -> index
+ return %lb : index
+}
+
// -----
// CHECK-LABEL: func @arith_divsi_negative_positive()
@@ -385,6 +404,62 @@ func.func @arith_select(%c: i1) -> (index, index) {
// -----
+// A merge with one candidate is an equality.
+// CHECK-LABEL: func @single_candidate_merge(
+func.func @single_candidate_merge() -> index {
+ // CHECK: arith.constant 7 : index
+ %c7 = arith.constant 7 : index
+ %r = test.merge_values %c7 : index
+ // CHECK: %[[eq:.*]] = arith.constant 7 : index
+ %eq = "test.reify_bound"(%r) {type = "EQ", constant} : (index) -> index
+ // CHECK: return %[[eq]]
+ return %eq : index
+}
+
+// -----
+
+// More than two candidates use the minimum and maximum candidates as a
+// conservative envelope. Candidate order is intentionally not sorted.
+// CHECK-LABEL: func @multi_candidate_merge(
+func.func @multi_candidate_merge() -> (index, index) {
+ // CHECK: arith.constant 11 : index
+ %c11 = arith.constant 11 : index
+ // CHECK: arith.constant 3 : index
+ %c3 = arith.constant 3 : index
+ // CHECK: arith.constant 7 : index
+ %c7 = arith.constant 7 : index
+ // CHECK: arith.constant 5 : index
+ %c5 = arith.constant 5 : index
+ %r = test.merge_values %c11, %c3, %c7, %c5 : index
+ // CHECK: %[[lb:.*]] = arith.constant 3 : index
+ // CHECK: %[[ub:.*]] = arith.constant 12 : index
+ %lb = "test.reify_bound"(%r) {type = "LB"} : (index) -> index
+ %ub = "test.reify_bound"(%r) {type = "UB"} : (index) -> index
+ // CHECK: return %[[lb]], %[[ub]]
+ return %lb, %ub : index, index
+}
+
+// -----
+
+// The result dimension and candidate dimensions are independent. Here result
+// dimension 0 merges lhs dimension 1 (size 4) and rhs dimension 0 (size 9).
+// CHECK-LABEL: func @merge_different_dims(
+func.func @merge_different_dims(
+ %lhs: tensor<?x4xf32>, %rhs: tensor<9x?xf32>) -> (index, index) {
+ %r = "test.merge_different_dims"(%lhs, %rhs)
+ : (tensor<?x4xf32>, tensor<9x?xf32>) -> tensor<?xf32>
+ // CHECK: %[[lb:.*]] = arith.constant 4 : index
+ // CHECK: %[[ub:.*]] = arith.constant 10 : index
+ %lb = "test.reify_bound"(%r) {dim = 0, type = "LB"}
+ : (tensor<?xf32>) -> index
+ %ub = "test.reify_bound"(%r) {dim = 0, type = "UB"}
+ : (tensor<?xf32>) -> index
+ // CHECK: return %[[lb]], %[[ub]]
+ return %lb, %ub : index, index
+}
+
+// -----
+
// CHECK-LABEL: func @arith_select_elementwise(
// CHECK-SAME: %[[a:.*]]: tensor<?xf32>, %[[b:.*]]: tensor<?xf32>, %[[c:.*]]: tensor<?xi1>)
func.func @arith_select_elementwise(%a: tensor<?xf32>, %b: tensor<?xf32>, %c: tensor<?xi1>) -> index {
diff --git a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
index 60fe96d52d20b..caabaa1010306 100644
--- a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir
@@ -192,6 +192,40 @@ func.func @scf_if_dynamic(%a: index, %b: index, %c : i1) {
// -----
+// A merge may depend on another merge. The common solver must revisit the
+// outer relationship after deriving the inner relationship's envelope.
+// CHECK-LABEL: func @scf_if_nested_merge(
+func.func @scf_if_nested_merge(%c1 : i1, %c2 : i1) {
+ // Consume the constants in the input IR before matching the newly reified
+ // lower and upper bounds below.
+ // CHECK: arith.constant 4 : index
+ // CHECK: arith.constant 9 : index
+ // CHECK: arith.constant 12 : index
+ %c4 = arith.constant 4 : index
+ %c9 = arith.constant 9 : index
+ %c12 = arith.constant 12 : index
+ %inner = scf.if %c1 -> index {
+ scf.yield %c4 : index
+ } else {
+ scf.yield %c9 : index
+ }
+ %outer = scf.if %c2 -> index {
+ scf.yield %inner : index
+ } else {
+ scf.yield %c12 : index
+ }
+
+ // CHECK: %[[c4:.*]] = arith.constant 4 : index
+ // CHECK: %[[c13:.*]] = arith.constant 13 : index
+ %lb = "test.reify_bound"(%outer) {type = "LB"} : (index) -> (index)
+ %ub = "test.reify_bound"(%outer) {type = "UB"} : (index) -> (index)
+ // CHECK: "test.some_use"(%[[c4]], %[[c13]])
+ "test.some_use"(%lb, %ub) : (index, index) -> ()
+ return
+}
+
+// -----
+
func.func @scf_if_no_affine_bound(%a: index, %b: index, %c : i1) {
%r = scf.if %c -> index {
scf.yield %a : index
diff --git a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
index f11c14bc09fb6..9cab2a6b33e36 100644
--- a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
@@ -112,18 +112,15 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
// Prepare stop condition. By default, reify in terms of the op's
// operands. No stop condition is used when a constant was requested.
- std::function<bool(Value, std::optional<int64_t>,
- ValueBoundsConstraintSet & cstr)>
- stopCondition = [&](Value v, std::optional<int64_t> d,
- ValueBoundsConstraintSet &cstr) {
+ ValueBoundsConstraintSet::StopConditionFn stopCondition =
+ [&](ValueDim valueDim, ValueBoundsConstraintSet &cstr) {
// Reify in terms of SSA values that are different from `value`.
- return v != value;
+ return valueDim.first != value;
};
if (reifyToFuncArgs) {
// Reify in terms of function block arguments.
- stopCondition = [](Value v, std::optional<int64_t> d,
- ValueBoundsConstraintSet &cstr) {
- auto bbArg = dyn_cast<BlockArgument>(v);
+ stopCondition = [](ValueDim valueDim, ValueBoundsConstraintSet &cstr) {
+ auto bbArg = dyn_cast<BlockArgument>(valueDim.first);
if (!bbArg)
return false;
return isa<FunctionOpInterface>(bbArg.getParentBlock()->getParentOp());
@@ -143,10 +140,10 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
options.closedUB = true;
auto reifiedScalable =
vector::ScalableValueBoundsConstraintSet::computeScalableBound(
- value, dim, *op.getVscaleMin(), *op.getVscaleMax(), boundType,
+ {value, dim}, *op.getVscaleMin(), *op.getVscaleMax(), boundType,
options);
if (succeeded(reifiedScalable)) {
- SmallVector<std::pair<Value, std::optional<int64_t>>, 1> vscaleOperand;
+ SmallVector<ValueDim, 1> vscaleOperand;
if (reifiedScalable->map.getNumInputs() == 1) {
// The only possible input to the bound is vscale.
vscaleOperand.push_back(std::make_pair(
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 9fe2168e11737..29e0c0727a88b 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1058,6 +1058,37 @@ void TestValueWithBoundsOp::populateBoundsForIndexValue(
cstr.bound(v) <= getMax().getSExtValue();
}
+//===----------------------------------------------------------------------===//
+// TestMergeValuesOp
+//===----------------------------------------------------------------------===//
+
+void TestMergeValuesOp::populateBoundsForIndexValue(
+ Value value, ValueBoundsConstraintSet &cstr) {
+ // Collect the candidates locally and publish the complete relationship in a
+ // single call. A temporary query can therefore never observe only a prefix
+ // of this list and mistake it for a one-candidate equality.
+ ValueDimList candidates;
+ for (Value candidate : getCandidates())
+ candidates.emplace_back(candidate, std::nullopt);
+ cstr.addMerge({value, std::nullopt}, std::move(candidates));
+}
+
+//===----------------------------------------------------------------------===//
+// TestMergeDifferentDimsOp
+//===----------------------------------------------------------------------===//
+
+void TestMergeDifferentDimsOp::populateBoundsForShapedValueDim(
+ Value value, int64_t dim, ValueBoundsConstraintSet &cstr) {
+ if (dim != 0)
+ return;
+
+ // The result side and each candidate side select their dimensions
+ // independently. In particular, the first edge relates dimension 0 to
+ // dimension 1; a single merge-wide dimension parameter could not express it.
+ ValueDimList candidates{{getLhs(), /*dim=*/1}, {getRhs(), /*dim=*/0}};
+ cstr.addMerge({value, /*dim=*/0}, std::move(candidates));
+}
+
//===----------------------------------------------------------------------===//
// ReifyBoundOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index c3e06dc968a27..f3f13284e73ca 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2741,6 +2741,32 @@ def TestValueWithBoundsOp : TEST_Op<"value_with_bounds", [
let assemblyFormat = "attr-dict";
}
+def TestMergeValuesOp : TEST_Op<"merge_values", [
+ DeclareOpInterfaceMethods<ValueBoundsOpInterface, ["populateBoundsForIndexValue"]>
+ ]> {
+ let description = [{
+ Creates an index value whose possible origins are the variadic `candidates`.
+ This operation exists to test atomic merge collection: its ValueBounds
+ implementation passes all operands to one `addMerge` call.
+ }];
+ let arguments = (ins Variadic<Index>:$candidates);
+ let results = (outs Index:$result);
+ let assemblyFormat = "$candidates attr-dict `:` type($result)";
+}
+
+def TestMergeDifferentDimsOp : TEST_Op<"merge_different_dims", [
+ DeclareOpInterfaceMethods<ValueBoundsOpInterface, ["populateBoundsForShapedValueDim"]>
+ ]> {
+ let description = [{
+ Test-only shaped merge whose result dimension 0 may originate from
+ `lhs` dimension 1 or `rhs` dimension 0. This intentionally relates
+ different dimensions to exercise the independent dimension parameters of
+ `addMerge`.
+ }];
+ let arguments = (ins AnyRankedTensor:$lhs, AnyRankedTensor:$rhs);
+ let results = (outs AnyRankedTensor:$result);
+}
+
def ReifyBoundOp : TEST_Op<"reify_bound", [Pure]> {
let description = [{
More information about the Mlir-commits
mailing list