[mlir] Extend affine.min/max ValueBoundsOpInterfaceImpls #118840
Conversation
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-affine Author: None (Max191) ChangesThe ValueBoundsOpInterface implementation for affine.min and affine.max only supported bounding on one side (LB for affine.max and UB for affine.min). However, it is possible to be smarter about the bounding in some cases, and bound values on both sides for these ops. An affine.min op result can be lower bounded if a lower bound can be computed for every result expression in its affine map. In this case, the new lower bound can be set as the minimum of the lower bounds of all result expressions. The converse can be done for affine.max, computing the maximum of all result expression bounds. This PR adds this logic to the ValueBoundsOpInterface implementations of affine.min and affine.max. Full diff: https://github.com/llvm/llvm-project/pull/118840.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
index 82a9fb0d490882..a575c0603aee70 100644
--- a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -67,6 +67,27 @@ struct AffineMinOpInterface
expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
cstr.bound(value) <= bound;
}
+ // Get all constant lower bounds, choose minimum, and set lower bound to it.
+ MLIRContext *ctx = op->getContext();
+ AffineMap map = minOp.getAffineMap();
+ SmallVector<Value> mapOperands = minOp.getOperands();
+ std::optional<int64_t> minBound;
+ for (AffineExpr expr : map.getResults()) {
+ auto exprMap =
+ AffineMap::get(map.getNumDims(), map.getNumSymbols(), expr, ctx);
+ ValueBoundsConstraintSet::Variable exprVar(exprMap, mapOperands);
+ FailureOr<int64_t> exprBound =
+ cstr.computeConstantBound(presburger::BoundType::LB, exprVar,
+ /*stopCondition=*/nullptr);
+ // If any LB cannot be computed, then the total LB cannot be known.
+ if (failed(exprBound))
+ return;
+ if (!minBound.has_value() || exprBound.value() < minBound.value())
+ minBound = exprBound.value();
+ }
+ if (!minBound.has_value())
+ return;
+ cstr.bound(value) >= minBound.value();
};
};
@@ -88,6 +109,27 @@ struct AffineMaxOpInterface
expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
cstr.bound(value) >= bound;
}
+ // Get all constant upper bounds, choose maximum, and set upper bound to it.
+ MLIRContext *ctx = op->getContext();
+ AffineMap map = maxOp.getAffineMap();
+ SmallVector<Value> mapOperands = maxOp.getOperands();
+ std::optional<int64_t> maxBound;
+ for (AffineExpr expr : map.getResults()) {
+ auto exprMap =
+ AffineMap::get(map.getNumDims(), map.getNumSymbols(), expr, ctx);
+ ValueBoundsConstraintSet::Variable exprVar(exprMap, mapOperands);
+ FailureOr<int64_t> exprBound = cstr.computeConstantBound(
+ presburger::BoundType::UB, exprVar,
+ /*stopCondition=*/nullptr, /*closedUB=*/true);
+ // If any UB cannot be computed, then the total UB cannot be known.
+ if (failed(exprBound))
+ return;
+ if (!maxBound.has_value() || exprBound.value() > maxBound.value())
+ maxBound = exprBound.value();
+ }
+ if (!maxBound.has_value())
+ return;
+ cstr.bound(value) <= maxBound.value();
};
};
diff --git a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
index 935c08aceff548..8014a9c3d087ae 100644
--- a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
@@ -38,6 +38,19 @@ func.func @affine_max_ub(%a: index) -> (index) {
// -----
+// CHECK-LABEL: func @affine_max_const_ub(
+// CHECK-SAME: %[[a:.*]]: index
+// CHECK: %[[c5:.*]] = arith.constant 5 : index
+// CHECK: return %[[c5]]
+func.func @affine_max_const_ub(%a: index) -> (index) {
+ %0 = affine.min affine_map<(d0) -> (d0, 4)>(%a)
+ %1 = affine.max affine_map<(d0) -> (d0, 2)>(%0)
+ %2 = "test.reify_bound"(%1) {type = "UB"}: (index) -> (index)
+ return %2 : index
+}
+
+// -----
+
// CHECK-LABEL: func @affine_min_ub(
// CHECK-SAME: %[[a:.*]]: index
// CHECK: %[[c3:.*]] = arith.constant 3 : index
@@ -61,6 +74,19 @@ func.func @affine_min_lb(%a: index) -> (index) {
// -----
+// CHECK-LABEL: func @affine_min_const_lb(
+// CHECK-SAME: %[[a:.*]]: index
+// CHECK: %[[c0:.*]] = arith.constant 0 : index
+// CHECK: return %[[c0]]
+func.func @affine_min_const_lb(%a: index) -> (index) {
+ %0 = affine.max affine_map<(d0) -> (d0, 0)>(%a)
+ %1 = affine.min affine_map<(d0) -> (d0, 2)>(%0)
+ %2 = "test.reify_bound"(%1) {type = "LB"}: (index) -> (index)
+ return %2 : index
+}
+
+// -----
+
// CHECK-LABEL: func @composed_affine_apply(
// CHECK: %[[cst:.*]] = arith.constant -8 : index
// CHECK: return %[[cst]]
|
MaheshRavishankar
left a comment
MaheshRavishankar
left a comment
There was a problem hiding this comment.
Looks fine to me, but I dont have all the context here. Maybe wait for @Groverkss or @matthias-springer to take a look.
Signed-off-by: Max Dawkins <[email protected]>
Max191
force-pushed
the
extend-affine-min-max-value-bounds
branch
from
a1ccd9e to
34e3e0c
Compare
|
Just rebased this PR and going to try to land again. Gentle ping @matthias-springer @Groverkss |
| AffineMap::get(map.getNumDims(), map.getNumSymbols(), expr, ctx); | ||
| ValueBoundsConstraintSet::Variable exprVar(exprMap, mapOperands); | ||
| FailureOr<int64_t> exprBound = | ||
| cstr.computeConstantBound(presburger::BoundType::LB, exprVar, |
There was a problem hiding this comment.
This is a bit misleading: You're actually calling a static function ValueBoundsConstraintSet::computeConstantBound here. That means building a brand new constraint set, which can be expensive because the same IR may be re-analyzed.
Can you try using cstr.populateAndCompare instead? And keep track of which AffineExpr result is the current minimum one. Keep comparing against that one to find the overall smallest one. That should also work with non-constant bounds then. And the stopFunction will be set correctly.
There was a problem hiding this comment.
You can also take a look at IfOpInterface::populateBounds, which does something similar.
There was a problem hiding this comment.
I see, thanks for the suggestion. I'll look into these options.
There was a problem hiding this comment.
@matthias-springer I tried the cstr.populateAndCompare approach, but I don't think it can achieve the same result. Consider the example from the test:
func.func @affine_min_const_lb(%a: index) -> (index) {
%0 = affine.max affine_map<(d0) -> (d0, 0)>(%a)
%1 = affine.min affine_map<(d0) -> (d0, 2)>(%0)
%2 = "test.reify_bound"(%1) {type = "LB"}: (index) -> (index)
return %2 : index
}
We want the LB of %2, so we try to find the minimum result expr of the affine.min map. Comparing d0 with 2 gives an uncertain result, because d0 is only constrained to be >= 0. This means d0 could be less than, equal to, or greater than 2, so we can't determine the min expr in this way. In fact, if we were able to determine a min expr this way, then the affine.min op could be folded away all together.
The difference in what this PR is trying to achieve is that it is computing the lower bounds of each expr separately, and then comparing the bounds to each other. So in this example, even though d0 could be less than or greater than 2, the lower bounds of each result expr can be shown to be greater than or equal to 0, and the lower bound of %2 is 0.
For what this PR wants to do, I think we need to be computing the full bound of each expr. I don't have a good simple solution that avoids extra computation right now, but having this additional analysis can open additional optimizations in some cases. To make this work without the extra computation, I think we would need to make the computeConstantBound function non-static, and somehow track state of the bounds of each Value in the ValueBoundsConstraintSet.
Do you have any thoughts on what the best option is?
There was a problem hiding this comment.
Let me make sure I understand the problem.
Can you double check that my math is correct here?
func.func @affine_min_const_lb(%a: index) -> (index) {
%0 = affine.max affine_map<(d0) -> (d0, 0)>(%a)
%1 = affine.min affine_map<(d0) -> (d0, 2)>(%0)
%2 = "test.reify_bound"(%1) {type = "LB"}: (index) -> (index)
return %2 : index
}
%0 >= %a
%0 >= 0
LB(%0) = {%a, 0} // both are lower bounds
%1 <= 2
%1 <= %0
LB(%1) = min(LB(2), LB(%0))
= min(2, {%a, 0})
= {min(2, %a), min(2, 0)}
= {min(2, %a), 0} // both are lower boundsWhen there are multiple LBs / UBs, the infrastructure may return either one of the two bounds. In practice, for affine.min/affine.max, it currently returns the constant bound if there is one (see FlatLinearConstraints::getSliceBounds for details).
The problem with populateAndCompare is that it compares two SSA values, whereas we have to compare the LBs of two SSA values here.
So we need something like populateAndComputeBound. This allows us to compute LBs for both affine.min operands in the same constraint set. These LBs can then be compared with the existing populateAndCompare.
populateAndComputeBound would be similar to ValueBoundsConstraintSet::computeBound, but reuses the existing constraint set. It's not a static function. Also, it does not project out any variables.
Can you explore this direction a bit? I think it's worth optimizing for performance here. I'm afraid, the current approach may have exponential runtime complexity in the worst case, as every time an affine.max/min op is visited, we start two new analyses for the remaining (not yet visited) IR.
There was a problem hiding this comment.
Yes, that sounds like exactly what we need! I can try this out, although I have some other things on my plate as well. I may start with a workaround solution downstream, and then come back to this when I have some more time, since it will take me some extra time to gain more context on the ValueBoundsOpInterface.
| cstr.bound(value) >= bound; | ||
| } | ||
| // Get all constant upper bounds, choose maximum, and set upper bound to it. | ||
| MLIRContext *ctx = op->getContext(); |
There was a problem hiding this comment.
This piece of code looks very similar to the one above. Maybe some functionality can be put into a shared static function. Not sure if it's really possible / practical...
|
Going to convert this PR to a draft for now, and try a downstream solution until I am ready to try implementing the solution proposed by Matthias. |
4 participants
The ValueBoundsOpInterface implementation for affine.min and affine.max only supported bounding on one side (LB for affine.max and UB for affine.min). However, it is possible to be smarter about the bounding in some cases, and bound values on both sides for these ops. An affine.min op result can be lower bounded if a lower bound can be computed for every result expression in its affine map. In this case, the new lower bound can be set as the minimum of the lower bounds of all result expressions. The converse can be done for affine.max, computing the maximum of all result expression bounds. This PR adds this logic to the ValueBoundsOpInterface implementations of affine.min and affine.max.