76435f2dca
This commit adds support for `scf.if` to `ValueBoundsConstraintSet`.
Example:
```
%0 = scf.if ... -> index {
scf.yield %a : index
} else {
scf.yield %b : index
}
```
The following constraints hold for %0:
* %0 >= min(%a, %b)
* %0 <= max(%a, %b)
Such constraints cannot be added to the constraint set; min/max is not
supported by `IntegerRelation`. However, if we know which one of %a and
%b is larger, we can add constraints for %0. E.g., if %a <= %b:
* %0 >= %a
* %0 <= %b
This commit required a few minor changes to the
`ValueBoundsConstraintSet` infrastructure, so that values can be
compared while we are still in the process of traversing the IR/adding
constraints.
Note: This is a re-upload of #85895, which was reverted. The bug that
caused the failure was fixed in #87859.
117 lines
4.0 KiB
C++
117 lines
4.0 KiB
C++
//===- ScalableValueBoundsConstraintSet.cpp - Scalable Value Bounds -------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
|
|
|
|
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
|
|
|
namespace mlir::vector {
|
|
|
|
FailureOr<ConstantOrScalableBound::BoundSize>
|
|
ConstantOrScalableBound::getSize() const {
|
|
if (map.isSingleConstant())
|
|
return BoundSize{map.getSingleConstantResult(), /*scalable=*/false};
|
|
if (map.getNumResults() != 1 || map.getNumInputs() != 1)
|
|
return failure();
|
|
auto binop = dyn_cast<AffineBinaryOpExpr>(map.getResult(0));
|
|
if (!binop || binop.getKind() != AffineExprKind::Mul)
|
|
return failure();
|
|
auto matchConstant = [&](AffineExpr expr, int64_t &constant) -> bool {
|
|
if (auto cst = dyn_cast<AffineConstantExpr>(expr)) {
|
|
constant = cst.getValue();
|
|
return true;
|
|
}
|
|
return false;
|
|
};
|
|
// Match `s0 * cst` or `cst * s0`:
|
|
int64_t cst = 0;
|
|
auto lhs = binop.getLHS();
|
|
auto rhs = binop.getRHS();
|
|
if ((matchConstant(lhs, cst) && isa<AffineSymbolExpr>(rhs)) ||
|
|
(matchConstant(rhs, cst) && isa<AffineSymbolExpr>(lhs))) {
|
|
return BoundSize{cst, /*scalable=*/true};
|
|
}
|
|
return failure();
|
|
}
|
|
|
|
char ScalableValueBoundsConstraintSet::ID = 0;
|
|
|
|
FailureOr<ConstantOrScalableBound>
|
|
ScalableValueBoundsConstraintSet::computeScalableBound(
|
|
Value value, std::optional<int64_t> dim, unsigned vscaleMin,
|
|
unsigned vscaleMax, presburger::BoundType boundType, bool closedUB,
|
|
StopConditionFn stopCondition) {
|
|
using namespace presburger;
|
|
assert(vscaleMin <= vscaleMax);
|
|
|
|
// No stop condition specified: Keep adding constraints until the worklist
|
|
// is empty.
|
|
auto defaultStopCondition = [&](Value v, std::optional<int64_t> dim,
|
|
mlir::ValueBoundsConstraintSet &cstr) {
|
|
return false;
|
|
};
|
|
|
|
ScalableValueBoundsConstraintSet scalableCstr(
|
|
value.getContext(), stopCondition ? stopCondition : defaultStopCondition,
|
|
vscaleMin, vscaleMax);
|
|
int64_t pos = scalableCstr.insert(value, dim, /*isSymbol=*/false);
|
|
scalableCstr.processWorklist();
|
|
|
|
// 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;
|
|
};
|
|
scalableCstr.projectOut(projectOutFn);
|
|
|
|
assert(scalableCstr.cstr.getNumDimAndSymbolVars() ==
|
|
scalableCstr.positionToValueDim.size() &&
|
|
"inconsistent mapping state");
|
|
|
|
// Check that the only columns left are vscale and the starting point.
|
|
for (int64_t i = 0; i < scalableCstr.cstr.getNumDimAndSymbolVars(); ++i) {
|
|
if (i == pos)
|
|
continue;
|
|
if (scalableCstr.positionToValueDim[i] !=
|
|
ValueDim(scalableCstr.getVscaleValue(),
|
|
ValueBoundsConstraintSet::kIndexValue)) {
|
|
return failure();
|
|
}
|
|
}
|
|
|
|
SmallVector<AffineMap, 1> lowerBound(1), upperBound(1);
|
|
scalableCstr.cstr.getSliceBounds(pos, 1, value.getContext(), &lowerBound,
|
|
&upperBound, closedUB);
|
|
|
|
auto invalidBound = [](auto &bound) {
|
|
return !bound[0] || bound[0].getNumResults() != 1;
|
|
};
|
|
|
|
AffineMap bound = [&] {
|
|
if (boundType == BoundType::EQ && !invalidBound(lowerBound) &&
|
|
lowerBound[0] == lowerBound[0]) {
|
|
return lowerBound[0];
|
|
} else if (boundType == BoundType::LB && !invalidBound(lowerBound)) {
|
|
return lowerBound[0];
|
|
} else if (boundType == BoundType::UB && !invalidBound(upperBound)) {
|
|
return upperBound[0];
|
|
}
|
|
return AffineMap{};
|
|
}();
|
|
|
|
if (!bound)
|
|
return failure();
|
|
|
|
return ConstantOrScalableBound{bound};
|
|
}
|
|
|
|
} // namespace mlir::vector
|