[mlir][Interfaces][NFC] ValueBoundsConstraintSet: Add columns for constant values/dims
#86097
Conversation
|
@llvm/pr-subscribers-mlir Author: Matthias Springer (matthias-springer) Changes
With this commit, dynamic and static values/dimension sizes are treated in the same way: in either case, a dimension/symbol is added to the constraint set. This is needed for a subsequent commit that adds support for branches. Full diff: https://github.com/llvm/llvm-project/pull/86097.diff 2 Files Affected:
diff --git a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
index 28dadfb9ecf868..94a8a8b429c801 100644
--- a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
+++ b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
@@ -280,12 +280,13 @@ class ValueBoundsConstraintSet {
/// 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
- /// worklist.
+ /// worklist if `addToWorklist` is set.
///
/// 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(Value value, std::optional<int64_t> dim, bool isSymbol = true,
+ bool addToWorklist = true);
/// Insert an anonymous column into the constraint set. The column is not
/// bound to any value/dimension. If `isSymbol` is set to "false", a dimension
diff --git a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
index 85abc2df894797..02af3a83166dfb 100644
--- a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
+++ b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
@@ -105,25 +105,57 @@ AffineExpr ValueBoundsConstraintSet::getExpr(Value value,
assertValidValueDim(value, dim);
#endif // NDEBUG
+ // Helper function that returns an affine expression that represents column
+ // `pos` in the constraint set.
+ auto getPosExpr = [&](int64_t pos) {
+ assert(pos >= 0 && pos < cstr.getNumDimAndSymbolVars() &&
+ "invalid position");
+ return pos < cstr.getNumDimVars()
+ ? builder.getAffineDimExpr(pos)
+ : builder.getAffineSymbolExpr(pos - cstr.getNumDimVars());
+ };
+
+ // Check if the value/dim is statically known. In that case, an affine
+ // constant expression should be returned. This allows us to support
+ // multiplications with constants. (Multiplications of two columns in the
+ // constraint set is not supported.)
+ std::optional<int64_t> constSize = std::nullopt;
auto shapedType = dyn_cast<ShapedType>(value.getType());
if (shapedType) {
- // Static dimension: return constant directly.
if (shapedType.hasRank() && !shapedType.isDynamicDim(*dim))
- return builder.getAffineConstantExpr(shapedType.getDimSize(*dim));
- } else {
- // Constant index value: return directly.
- if (auto constInt = ::getConstantIntValue(value))
- return builder.getAffineConstantExpr(*constInt);
+ constSize = shapedType.getDimSize(*dim);
+ } else if (auto constInt = ::getConstantIntValue(value)) {
+ constSize = *constInt;
}
- // Dynamic value: add to constraint set.
+ // 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))
- (void)insert(value, dim);
- int64_t pos = getPos(value, dim);
- return pos < cstr.getNumDimVars()
- ? builder.getAffineDimExpr(pos)
- : builder.getAffineSymbolExpr(pos - cstr.getNumDimVars());
+ 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));
+ }
+
+ 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);
+ if (shapedType)
+ bound(value)[*dim] == *constSize;
+ else
+ bound(value) == *constSize;
+ return builder.getAffineConstantExpr(*constSize);
+ }
+
+ // 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));
}
AffineExpr ValueBoundsConstraintSet::getExpr(OpFoldResult ofr) {
@@ -140,7 +172,7 @@ AffineExpr ValueBoundsConstraintSet::getExpr(int64_t constant) {
int64_t ValueBoundsConstraintSet::insert(Value value,
std::optional<int64_t> dim,
- bool isSymbol) {
+ bool isSymbol, bool addToWorklist) {
#ifndef NDEBUG
assertValidValueDim(value, dim);
#endif // NDEBUG
@@ -155,7 +187,12 @@ int64_t ValueBoundsConstraintSet::insert(Value value,
if (positionToValueDim[i].has_value())
valueDimToPosition[*positionToValueDim[i]] = i;
- worklist.push(pos);
+ if (addToWorklist) {
+ LLVM_DEBUG(llvm::dbgs() << "Push to worklist: " << value
+ << " (dim: " << dim.value_or(kIndexValue) << ")\n");
+ worklist.push(pos);
+ }
+
return pos;
}
|
matthias-springer
force-pushed
the
users/matthias-springer/value_bounds_const_columns
branch
from
13c9d7f to
6bbcba6
Compare
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
|
✅ With the latest revision this PR passed the Python code formatter. |
matthias-springer
force-pushed
the
users/matthias-springer/value_bounds_const_columns
branch
from
6bbcba6 to
d53abc4
Compare
| // Constant index value: return directly. | ||
| if (auto constInt = ::getConstantIntValue(value)) | ||
| return builder.getAffineConstantExpr(*constInt); | ||
| constSize = shapedType.getDimSize(*dim); |
There was a problem hiding this comment.
Hmm, I feel this is tricky: this really depends about the context in which we query the analysis.
It is easy to create a varying shape (i.e. size depends on loop iv) and this will not be a symbol anymore.
Basically, if the analysis was properly rooted at an op / block, we could say that everything defined above is a symbol.
IOW, any such analysis is relative to its context: if your transform/analysis cares about information on say k innermost loops, you can "make" information about the outermost loops constant.
There was a problem hiding this comment.
I don't fully understand the different between symbols and dimensions in IntegerRelation. The limitations that I am aware of:
getSliceBoundscan only be called on a dimensions, so the first "column" that I add to the constraint set (for which we want to compute LB/UB) is a dimension.projectOutworks on both symbols and dimensions, so all other values are added as symbols.
Basically we have only one dimension and everything else is a symbol. (But the current implementation would make it easy to change this.) This seems to work well so far.
There was a problem hiding this comment.
any such analysis is relative to its context
On second thought, that's actually how it is implemented at the moment. We provide a ValueBoundsConstraintSet::computeBound(Value) API. That function will build a constraint set and add the given value as a dimension, everything else as a symbol. When computing a bound for a different value, we build a brand new constraint set, this time with the other value being the only dimension. (There is no public API that exposes the constraint set itself.) So this analysis is actually relative to a context and rooted at a certain op.
Computing multiple bounds back-to-back is slow at the moment because we have to re-analyze the entire IR. If we improve this at some point in the future, we will have to take a closer look at what's a symbol and what's a dimension, as you suggested.
Yes, I think so.
A few more constraints are added, but not many. E.g., when we have IR such as: %0 = tensor.insert %f into %t[%idx] : tensor<5xf32>
%1 = tensor.insert %f into %0[%idx] : tensor<5xf32>
%2 = tensor.insert %f into %1[%idx] : tensor<5xf32>
"test.reify_bound"(%2) {dim = 0}Only one constraint will be added: There's no easy way to test this. The information in the constraint system does not change, so the computed bounds are exactly the same. That's why I marked it NFC. This change is needed for #85895 and I put it into an extra PR to keep the other PRs smaller. ( |
|
I could make this PR part of #85895 and merge the two as one commit. That's where this change is needed, and an assertion would be triggered in one of the test cases without this change. |
…onstant values/dims `ValueBoundsConstraintSet` maintains an internal constraint set (`IntegerRelation`), where every analyzed index-typed SSA value or dimension of a shaped type is represented with a dimension/symbol. Prior to this change, index-typed values with a statically known constant value and static shaped type dimensions were not added to the constraint set. Instead, `getExpr` directly returned an affine constrant expression. With this commit, dynamic and static values/dimension sizes are treated in the same way: in either case, a dimension/symbol is added to the constraint set. This is needed for a subsequent commit that adds support for branches.
matthias-springer
force-pushed
the
users/matthias-springer/value_bounds_const_columns
branch
from
d53abc4 to
fca8ef5
Compare
|
I merged this PR into #85895. That's where it was originally and then we have a test that exercises the changes in the code. |
ValueBoundsConstraintSetmaintains an internal constraint set (IntegerRelation), where every analyzed index-typed SSA value or dimension of a shaped type is represented with a dimension/symbol. Prior to this change, index-typed values with a statically known constant value and static shaped type dimensions were not added to the constraint set. Instead,getExprdirectly returned an affine constant expression.With this commit, dynamic and static values/dimension sizes are treated in the same way: in either case, a dimension/symbol is added to the constraint set. This is needed for a subsequent commit that adds support for branches.