-
Notifications
You must be signed in to change notification settings - Fork 14.3k
[mlir][affine][Analysis] Add conservative bounds for semi-affine mods #93576
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
@llvm/pr-subscribers-mlir-presburger @llvm/pr-subscribers-mlir-core Author: Benjamin Maxwell (MacDue) ChangesThis 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:
There are caveats to this implementation. To be able to add a bound for a This is not a problem for computing scalable bounds where it's safe to assume 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:
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]
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the patch. The approach for approximating is good, I'm mostly concerned about layering.
I have a few concerns about the implementation adding these overapproximation very deep into Presburger/Affine APIs. I would expect this patch to not touch them at all. This is how I think this should look like:
- Add a new method to FlatLinearValueConstraints allowing adding bounds on semi-affine expressions with a defined way of approximating.
- This method would need a new flattened. IIUC, the AffineExprFlattener is designed to be derived. I would create a SemiAffineExprFlattener which adds these approximations and flattens them and use that in this method.
@@ -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, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: I would prefer if these two APIs looked more consistent. Could you give the members the same names and same orders? Also, please update docs for addLocalIdSemiAffine.
bc5f1c9
to
f295789
Compare
Just looking in the affine sources I found:
So does that mean if the (symbolic) rhs is negative for a semi-affine mod it's UB? |
This is not really my area of expertise, but looks like all comments from @Groverkss have been addressed. Kunwar, can you take another look? This LGTM :)
Do you have a link? Not supporting it might be safer - it's not obvious to me why it would be UB TBH. |
https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/AffineExpr.h#L45-L46. And I mean UB according to the affine dialect. If it's specified that the symbolic RHS is always positive, then it's not really defined for the negative case (which is still possible for symbolic expressions). If the negative case is not defined/intended to be supported that makes things easier for us (since our assumption that the |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm happy with the implementation, looks great, thanks! I just have a couple of documentation comments. Feel free to land once you address them.
virtual void addLocalIdSemiAffine(AffineExpr localExpr); | ||
virtual LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs, | ||
ArrayRef<int64_t> rhs, | ||
AffineExpr localExpr); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: Can you rename lhs -> dividend, rhs -> divisor (just to keep consistency with addLocalFloorDivId)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These won't always be a dividend
and divisor
. This is also called for semi-affine multiplication expressions, which is why these are just lhs
and rhs
.
mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
Outdated
Show resolved
Hide resolved
@@ -62,6 +60,10 @@ ScalableValueBoundsConstraintSet::computeScalableBound( | |||
int64_t pos = scalableCstr.insert(value, dim, /*isSymbol=*/false); | |||
scalableCstr.processWorklist(); | |||
|
|||
// Check the resulting constraints set is valid. | |||
if (scalableCstr.cstr.isEmpty()) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You should also check if isEmpty
is the solver you want to use. There is a Simplex solver which might be better as you are going into more complex constraints (or Integer solver if you have mod/floor constraints). Have a look at this PR for more context: #93010
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think .isEmpty()
does the job right now (I was mainly going off the existing value bounds usage). I think this would be fine to replace with .isEmptyByFMTest()
once your PR lands.
mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
Outdated
Show resolved
Hide resolved
If it's in AffineExpr definition, I would say this is part of AffineExpr, not affine dialect. I think it's ok to assume it and link it. If it is later removed from it's definition, we can remove this constraint (and add a parameter to specify if rhs is positive) |
This path 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).
Thanks for the review(s)! Unless there's any last minute comments I'll land this tomorrow 😄 |
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:
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 themod
expression. So to handle this a constraint is added forrhs > 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).