[Mlir-commits] [mlir] [mlir][affine][Analysis] Add conservative bounds for semi-affine mods (PR #93576)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue May 28 09:45:13 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-core
Author: Benjamin Maxwell (MacDue)
<details>
<summary>Changes</summary>
This patch adds support for computing bounds for semi-affine mod expression to FlatLinearConstraints. This is then enabled within the ScalableValueBoundsConstraintSet to allow computing the bounds of scalable remainder loops.
E.g. computing the bound of something like:
```
%0 = affine.apply #remainder_start_index()[%c8_vscale]
scf.for %i = %0 to %c1000 step %c8_vscale {
%remaining_iterations = affine.apply #remaining_iterations(%i)
// The upper bound for the remainder loop iterations should be:
// %c8_vscale - 1 (expressed as an affine map,
// affine_map<()[s0] -> (s0 * 8 - 1)>, where s0 is vscale)
%bound = "test.reify_bound"(%remaining_iterations) <{scalable, ...}>
}
```
There are caveats to this implementation. To be able to add a bound for a `mod` we need to assume the rhs is positive (> 0). This may not be known when adding the bounds for the `mod` expression. So to handle this a constraint is added for `rhs > 0`, this may later be found not to hold (in which case the constraints set becomes empty/invalid).
This is not a problem for computing scalable bounds where it's safe to assume `s0` is vscale (or some positive multiple of it). But this may need to be considered when enabling this feature elsewhere (to ensure correctness).
---
Patch is 32.50 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/93576.diff
11 Files Affected:
- (modified) mlir/include/mlir/Analysis/FlatLinearValueConstraints.h (+37-12)
- (modified) mlir/include/mlir/Analysis/Presburger/IntegerRelation.h (+14)
- (modified) mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h (+3-2)
- (modified) mlir/include/mlir/IR/AffineExprVisitor.h (+4-2)
- (modified) mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h (+5-1)
- (modified) mlir/lib/Analysis/FlatLinearValueConstraints.cpp (+70-31)
- (modified) mlir/lib/Analysis/Presburger/IntegerRelation.cpp (+31)
- (modified) mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp (+10)
- (modified) mlir/lib/IR/AffineExpr.cpp (+13-7)
- (modified) mlir/lib/Interfaces/ValueBoundsOpInterface.cpp (+14-4)
- (modified) mlir/test/Dialect/Vector/test-scalable-bounds.mlir (+56)
``````````diff
diff --git a/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h b/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
index 29c19442a7c7c..a85e2790373bd 100644
--- a/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
+++ b/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
@@ -66,6 +66,10 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
/// Return the kind of this object.
Kind getKind() const override { return Kind::FlatLinearConstraints; }
+ /// Flag to control if conservative semi-affine bounds should be added in
+ /// `addBound()`.
+ enum class AddConservativeSemiAffineBounds { No = 0, Yes };
+
/// Adds a bound for the variable at the specified position with constraints
/// being drawn from the specified bound map. In case of an EQ bound, the
/// bound map is expected to have exactly one result. In case of a LB/UB, the
@@ -77,21 +81,39 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
/// as a closed bound by +1/-1 respectively. In case of an EQ bound, it can
/// only be added as a closed bound.
///
+ /// Conservative bounds for semi-affine expressions will be added if
+ /// `AddConservativeSemiAffineBounds` is set to `Yes`. This currently does not
+ /// cover all semi-affine expressions, so `addBound()` still may fail with
+ /// this set. Note: If enabled it is possible for the resulting constraint set
+ /// to become empty if a precondition of a conservative bound is found not to
+ /// hold.
+ ///
/// Note: The dimensions/symbols of this FlatLinearConstraints must match the
/// dimensions/symbols of the affine map.
- LogicalResult addBound(presburger::BoundType type, unsigned pos,
- AffineMap boundMap, bool isClosedBound);
+ LogicalResult addBound(
+ presburger::BoundType type, unsigned pos, AffineMap boundMap,
+ bool isClosedBound,
+ AddConservativeSemiAffineBounds = AddConservativeSemiAffineBounds::No);
/// Adds a bound for the variable at the specified position with constraints
/// being drawn from the specified bound map. In case of an EQ bound, the
/// bound map is expected to have exactly one result. In case of a LB/UB, the
/// bound map may have more than one result, for each of which an inequality
/// is added.
+ ///
+ /// Conservative bounds for semi-affine expressions will be added if
+ /// `AddConservativeSemiAffineBounds` is set to `Yes`. This currently does not
+ /// cover all semi-affine expressions, so `addBound()` still may fail with
+ /// this set. If enabled it is possible for the resulting constraint set
+ /// to become empty if a precondition of a conservative bound is found not to
+ /// hold.
+ ///
/// Note: The dimensions/symbols of this FlatLinearConstraints must match the
/// dimensions/symbols of the affine map. By default the lower bound is closed
/// and the upper bound is open.
- LogicalResult addBound(presburger::BoundType type, unsigned pos,
- AffineMap boundMap);
+ LogicalResult addBound(
+ presburger::BoundType type, unsigned pos, AffineMap boundMap,
+ AddConservativeSemiAffineBounds = AddConservativeSemiAffineBounds::No);
/// The `addBound` overload above hides the inherited overloads by default, so
/// we explicitly introduce them here.
@@ -193,7 +215,8 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
/// Note: This is a shared helper function of `addLowerOrUpperBound` and
/// `composeMatchingMap`.
LogicalResult flattenAlignedMapAndMergeLocals(
- AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs);
+ AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
+ bool addConservativeSemiAffineBounds = false);
/// Prints the number of constraints, dimensions, symbols and locals in the
/// FlatLinearConstraints. Also, prints for each variable whether there is
@@ -468,18 +491,19 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
/// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the
/// dimensions, symbols, and additional variables that represent floor divisions
/// of dimensions, symbols, and in turn other floor divisions. Returns failure
-/// if 'expr' could not be flattened (i.e., semi-affine is not yet handled).
+/// if 'expr' could not be flattened (i.e., an unhandled semi-affine was found).
/// 'cst' contains constraints that connect newly introduced local variables
/// to existing dimensional and symbolic variables. See documentation for
/// AffineExprFlattener on how mod's and div's are flattened.
-LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
- unsigned numSymbols,
- SmallVectorImpl<int64_t> *flattenedExpr,
- FlatLinearConstraints *cst = nullptr);
+LogicalResult
+getFlattenedAffineExpr(AffineExpr expr, unsigned numDims, unsigned numSymbols,
+ SmallVectorImpl<int64_t> *flattenedExpr,
+ FlatLinearConstraints *cst = nullptr,
+ bool addConservativeSemiAffineBounds = false);
/// Flattens the result expressions of the map to their corresponding flattened
/// forms and set in 'flattenedExprs'. Returns failure if any expression in the
-/// map could not be flattened (i.e., semi-affine is not yet handled). 'cst'
+/// map could not be flattened (i.e., an unhandled semi-affine was found). 'cst'
/// contains constraints that connect newly introduced local variables to
/// existing dimensional and / symbolic variables. See documentation for
/// AffineExprFlattener on how mod's and div's are flattened. For all affine
@@ -490,7 +514,8 @@ LogicalResult getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
LogicalResult
getFlattenedAffineExprs(AffineMap map,
std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
- FlatLinearConstraints *cst = nullptr);
+ FlatLinearConstraints *cst = nullptr,
+ bool addConservativeSemiAffineBounds = false);
LogicalResult
getFlattenedAffineExprs(IntegerSet set,
std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index 163f365c623d7..c7e2e55372324 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -454,6 +454,20 @@ class IntegerRelation {
addLocalFloorDiv(getMPIntVec(dividend), MPInt(divisor));
}
+ /// Adds a new local variable as the mod of an affine function of other
+ /// variables. The coefficients of the operands of the mod are provided in
+ /// `lhs` and `rhs` respectively. Three constraints are added to provide a
+ /// conservative bound for the mod:
+ /// 1. rhs > 0 (assumption/precondition)
+ /// 2. lhs % rhs < rhs
+ /// 3. lhs % rhs >= 0
+ /// We ensure the rhs is positive so we can assume the result is positive.
+ void addLocalModConservativeBounds(ArrayRef<MPInt> lhs, ArrayRef<MPInt> rhs);
+ void addLocalModConservativeBounds(ArrayRef<int64_t> lhs,
+ ArrayRef<int64_t> rhs) {
+ addLocalModConservativeBounds(getMPIntVec(lhs), getMPIntVec(rhs));
+ }
+
/// Projects out (aka eliminates) `num` variables starting at position
/// `pos`. The resulting constraint system is the shadow along the dimensions
/// that still exist. This method may not always be integer exact.
diff --git a/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h b/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
index 67a6581eb2fb4..d682ef1131152 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
+++ b/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
@@ -33,8 +33,9 @@ struct ScalableValueBoundsConstraintSet
MLIRContext *context,
ValueBoundsConstraintSet::StopConditionFn stopCondition,
unsigned vscaleMin, unsigned vscaleMax)
- : RTTIExtends(context, stopCondition), vscaleMin(vscaleMin),
- vscaleMax(vscaleMax) {};
+ : RTTIExtends(context, stopCondition,
+ /*addConservativeSemiAffineBounds=*/true),
+ vscaleMin(vscaleMin), vscaleMax(vscaleMax){};
using RTTIExtends::bound;
using RTTIExtends::StopConditionFn;
diff --git a/mlir/include/mlir/IR/AffineExprVisitor.h b/mlir/include/mlir/IR/AffineExprVisitor.h
index 27c49cd80018e..bff9c9d4a029c 100644
--- a/mlir/include/mlir/IR/AffineExprVisitor.h
+++ b/mlir/include/mlir/IR/AffineExprVisitor.h
@@ -413,7 +413,8 @@ class SimpleAffineExprFlattener
/// lhs of the mod, floordiv, ceildiv or mul expression and with respect to a
/// symbolic rhs expression. `localExpr` is the simplified tree expression
/// (AffineExpr) corresponding to the quantifier.
- virtual void addLocalIdSemiAffine(AffineExpr localExpr);
+ virtual void addLocalIdSemiAffine(AffineExpr localExpr, ArrayRef<int64_t> lhs,
+ ArrayRef<int64_t> rhs);
private:
/// Adds `expr`, which may be mod, ceildiv, floordiv or mod expression
@@ -422,7 +423,8 @@ class SimpleAffineExprFlattener
/// quantifier is already present, we put the coefficient in the proper index
/// of `result`, otherwise we add a new local variable and put the coefficient
/// there.
- void addLocalVariableSemiAffine(AffineExpr expr,
+ void addLocalVariableSemiAffine(AffineExpr expr, ArrayRef<int64_t> lhs,
+ ArrayRef<int64_t> rhs,
SmallVectorImpl<int64_t> &result,
unsigned long resultSize);
diff --git a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
index ac17ace5a976d..337314143c80c 100644
--- a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
+++ b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
@@ -313,7 +313,8 @@ class ValueBoundsConstraintSet
/// An index-typed value or the dimension of a shaped-type value.
using ValueDim = std::pair<Value, int64_t>;
- ValueBoundsConstraintSet(MLIRContext *ctx, StopConditionFn stopCondition);
+ ValueBoundsConstraintSet(MLIRContext *ctx, StopConditionFn stopCondition,
+ bool addConservativeSemiAffineBounds = false);
/// Return "true" if, based on the current state of the constraint system,
/// "lhs cmp rhs" was proven to hold. Return "false" if the specified relation
@@ -404,6 +405,9 @@ class ValueBoundsConstraintSet
/// The current stop condition function.
StopConditionFn stopCondition = nullptr;
+
+ /// Should conservative bounds be added for semi-affine expressions.
+ bool addConservativeSemiAffineBounds = false;
};
} // namespace mlir
diff --git a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
index 8b38016d61498..b19aa2a06ba62 100644
--- a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
+++ b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
@@ -46,9 +46,15 @@ struct AffineExprFlattener : public SimpleAffineExprFlattener {
// inequalities.
IntegerPolyhedron localVarCst;
- AffineExprFlattener(unsigned nDims, unsigned nSymbols)
+ AffineExprFlattener(unsigned nDims, unsigned nSymbols,
+ bool addConservativeSemiAffineBounds = false)
: SimpleAffineExprFlattener(nDims, nSymbols),
- localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {}
+ localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)),
+ addConservativeSemiAffineBounds(addConservativeSemiAffineBounds) {}
+
+ bool hasUnhandledSemiAffineExpressions() const {
+ return unhandledSemiAffineExpressions;
+ }
private:
// Add a local variable (needed to flatten a mod, floordiv, ceildiv expr).
@@ -63,35 +69,61 @@ struct AffineExprFlattener : public SimpleAffineExprFlattener {
// Update localVarCst.
localVarCst.addLocalFloorDiv(dividend, divisor);
}
+
+ // Add a local identifier (needed to flatten a mod, floordiv, ceildiv, mul
+ // expr) when the rhs is a symbolic expression. The local identifier added
+ // may be a floordiv, ceildiv, mul or mod of a pure affine/semi-affine
+ // function of other identifiers, coefficients of which are specified in the
+ // lhs of the mod, floordiv, ceildiv or mul expression and with respect to a
+ // symbolic rhs expression. `localExpr` is the simplified tree expression
+ // (AffineExpr) corresponding to the quantifier.
+ void addLocalIdSemiAffine(AffineExpr localExpr, ArrayRef<int64_t> lhs,
+ ArrayRef<int64_t> rhs) override {
+ SimpleAffineExprFlattener::addLocalIdSemiAffine(localExpr, lhs, rhs);
+ if (!addConservativeSemiAffineBounds) {
+ unhandledSemiAffineExpressions = true;
+ return;
+ }
+ if (localExpr.getKind() == AffineExprKind::Mod) {
+ localVarCst.addLocalModConservativeBounds(lhs, rhs);
+ return;
+ }
+ // TODO: Support other semi-affine expressions.
+ unhandledSemiAffineExpressions = true;
+ }
+
+ bool addConservativeSemiAffineBounds = false;
+ bool unhandledSemiAffineExpressions = false;
};
} // namespace
// Flattens the expressions in map. Returns failure if 'expr' was unable to be
// flattened. For example two specific cases:
-// 1. semi-affine expressions not handled yet.
+// 1. an unhandled semi-affine expressions is found.
// 2. has poison expression (i.e., division by zero).
static LogicalResult
getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
unsigned numSymbols,
std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
- FlatLinearConstraints *localVarCst) {
+ FlatLinearConstraints *localVarCst,
+ bool addConservativeSemiAffineBounds = false) {
if (exprs.empty()) {
if (localVarCst)
*localVarCst = FlatLinearConstraints(numDims, numSymbols);
return success();
}
- AffineExprFlattener flattener(numDims, numSymbols);
+ AffineExprFlattener flattener(numDims, numSymbols,
+ addConservativeSemiAffineBounds);
// Use the same flattener to simplify each expression successively. This way
// local variables / expressions are shared.
for (auto expr : exprs) {
- if (!expr.isPureAffine())
- return failure();
- // has poison expression
auto flattenResult = flattener.walkPostOrder(expr);
if (failed(flattenResult))
return failure();
+ if (flattener.hasUnhandledSemiAffineExpressions())
+ return failure();
}
assert(flattener.operandExprStack.size() == exprs.size());
@@ -106,33 +138,33 @@ getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
}
// Flattens 'expr' into 'flattenedExpr'. Returns failure if 'expr' was unable to
-// be flattened (semi-affine expressions not handled yet).
-LogicalResult
-mlir::getFlattenedAffineExpr(AffineExpr expr, unsigned numDims,
- unsigned numSymbols,
- SmallVectorImpl<int64_t> *flattenedExpr,
- FlatLinearConstraints *localVarCst) {
+// be flattened (an unhandled semi-affine was found).
+LogicalResult mlir::getFlattenedAffineExpr(
+ AffineExpr expr, unsigned numDims, unsigned numSymbols,
+ SmallVectorImpl<int64_t> *flattenedExpr, FlatLinearConstraints *localVarCst,
+ bool addConservativeSemiAffineBounds) {
std::vector<SmallVector<int64_t, 8>> flattenedExprs;
- LogicalResult ret = ::getFlattenedAffineExprs({expr}, numDims, numSymbols,
- &flattenedExprs, localVarCst);
+ LogicalResult ret =
+ ::getFlattenedAffineExprs({expr}, numDims, numSymbols, &flattenedExprs,
+ localVarCst, addConservativeSemiAffineBounds);
*flattenedExpr = flattenedExprs[0];
return ret;
}
/// Flattens the expressions in map. Returns failure if 'expr' was unable to be
-/// flattened (i.e., semi-affine expressions not handled yet).
+/// flattened (i.e., an unhandled semi-affine was found).
LogicalResult mlir::getFlattenedAffineExprs(
AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
- FlatLinearConstraints *localVarCst) {
+ FlatLinearConstraints *localVarCst, bool addConservativeSemiAffineBounds) {
if (map.getNumResults() == 0) {
if (localVarCst)
*localVarCst =
FlatLinearConstraints(map.getNumDims(), map.getNumSymbols());
return success();
}
- return ::getFlattenedAffineExprs(map.getResults(), map.getNumDims(),
- map.getNumSymbols(), flattenedExprs,
- localVarCst);
+ return ::getFlattenedAffineExprs(
+ map.getResults(), map.getNumDims(), map.getNumSymbols(), flattenedExprs,
+ localVarCst, addConservativeSemiAffineBounds);
}
LogicalResult mlir::getFlattenedAffineExprs(
@@ -641,9 +673,11 @@ void FlatLinearConstraints::getSliceBounds(unsigned offset, unsigned num,
}
LogicalResult FlatLinearConstraints::flattenAlignedMapAndMergeLocals(
- AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs) {
+ AffineMap map, std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
+ bool addConservativeSemiAffineBounds) {
FlatLinearConstraints localCst;
- if (failed(getFlattenedAffineExprs(map, flattenedExprs, &localCst))) {
+ if (failed(getFlattenedAffineExprs(map, flattenedExprs, &localCst,
+ addConservativeSemiAffineBounds))) {
LLVM_DEBUG(llvm::dbgs()
<< "composition unimplemented for semi-affine maps\n");
return failure();
@@ -664,9 +698,9 @@ LogicalResult FlatLinearConstraints::flattenAlignedMapAndMergeLocals(
return success();
}
-LogicalResult FlatLinearConstraints::addBound(BoundType type, unsigned pos,
- AffineMap boundMap,
- bool isClosedBound) {
+LogicalResult FlatLinearConstraints::addBound(
+ BoundType type, unsigned pos, AffineMap boundMap, bool isClosedBound,
+ AddConservativeSemiAffineBounds addSemiAffineBounds) {
assert(boundMap.getNumDims() == getNumDimVars() && "dim mismatch");
assert(boundMap.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");
assert(pos < getNumDimAndSymbolVars() && "invalid position");
@@ -680,7 +714,9 @@ LogicalResult FlatLinearConstraints::addBound(BoundType type, unsigned pos,
bool lower = type == BoundType::LB || type == BoundType::EQ;
std::vector<SmallVector<int64_t, 8>> flatExprs;
- if (failed(flattenAlignedMapAndMergeLocals(boundMap, &flatExprs)))
+ if (failed(flattenAlignedMapAndMergeLocals(
+ boundMap, &flatExprs,
+ addSemiAffineBounds == AddConservativeSemiAffineBounds::Yes)))
return failure();
assert(flatExprs.size() == boundMap.getNumResults());
@@ -716,9 +752,11 @@ LogicalResult FlatLinearConstraints::addBound(BoundType type, unsigned pos,
return success();
}
-LogicalResult FlatLinearConstraints::addBound(BoundType type, unsigned pos,
- AffineMap boundMap) {
- return addBound(type, pos, boundMap, /*isClosedBound=*/type != BoundType::UB);
+LogicalResult FlatLinearConstraints::addBound(
+ BoundType type, unsigned pos, AffineMap boundMap,
+ AddConservativeSemiAffineBounds addSemiAffineBounds) {
+ return addBound(type, pos, boundMap,
+ /*isClosedBound=*/type != BoundType::UB, addSemiAffineBounds);
}
/// Compute an explicit representation for local vars. For all systems coming
@@ -1243,7 +1281,8 @@ mlir::getMultiAffineFunctionFromMap(AffineMap map,
"AffineMap cannot produce divs without local representation");
// TODO: We shouldn't have to do this conversion.
- Matrix<MPInt> mat(map.getNumResults(), map.getNumInputs() + divs.getNumDivs() + 1);
+ Matrix<MPInt> mat(map.getNumResults(),
+ map.getNumInputs() + divs.getNumDivs() + 1);
for (unsigned i = 0, e = ...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/93576
More information about the Mlir-commits
mailing list