[Mlir-commits] [mlir] 29a925a - [mlir][affine][Analysis] Add conservative bounds for semi-affine mods (#93576)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 5 03:35:17 PDT 2024
Author: Benjamin Maxwell
Date: 2024-06-05T11:35:13+01:00
New Revision: 29a925abb660104b413b15075b3a19793825f57e
URL: https://github.com/llvm/llvm-project/commit/29a925abb660104b413b15075b3a19793825f57e
DIFF: https://github.com/llvm/llvm-project/commit/29a925abb660104b413b15075b3a19793825f57e.diff
LOG: [mlir][affine][Analysis] Add conservative bounds for semi-affine mods (#93576)
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:
```
// `1000 mod s0` is a semi-affine.
#remainder_start_index = affine_map<()[s0] -> (-(1000 mod s0) + 1000)>
#remaining_iterations = affine_map<(d0) -> (-d0 + 1000)>
%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).
Added:
Modified:
mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
mlir/include/mlir/IR/AffineExprVisitor.h
mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
mlir/lib/Analysis/FlatLinearValueConstraints.cpp
mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
mlir/lib/IR/AffineExpr.cpp
mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
mlir/test/Dialect/Vector/test-scalable-bounds.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h b/mlir/include/mlir/Analysis/FlatLinearValueConstraints.h
index 29c19442a7c7c..cc6ab64b4b7d7 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 only
+ /// covers semi-affine `mod` expressions, so `addBound()` will still fail if
+ /// it encounters a semi-affine `floordiv`, `ceildiv`, or `mul`. 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 only
+ /// covers semi-affine `mod` expressions, so `addBound()` will still fail if
+ /// it encounters a semi-affine `floordiv`, `ceildiv`, or `mul`. 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. 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/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h b/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
index 67a6581eb2fb4..93b3c92533c54 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..fc4cd915d8453 100644
--- a/mlir/include/mlir/IR/AffineExprVisitor.h
+++ b/mlir/include/mlir/IR/AffineExprVisitor.h
@@ -413,18 +413,22 @@ 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 LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs,
+ ArrayRef<int64_t> rhs,
+ AffineExpr localExpr);
private:
- /// Adds `expr`, which may be mod, ceildiv, floordiv or mod expression
+ /// Adds `localExpr`, which may be mod, ceildiv, floordiv or mod expression
/// representing the affine expression corresponding to the quantifier
- /// introduced as the local variable corresponding to `expr`. If the
+ /// introduced as the local variable corresponding to `localExpr`. If the
/// 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,
- SmallVectorImpl<int64_t> &result,
- unsigned long resultSize);
+ LogicalResult addLocalVariableSemiAffine(ArrayRef<int64_t> lhs,
+ ArrayRef<int64_t> rhs,
+ AffineExpr localExpr,
+ SmallVectorImpl<int64_t> &result,
+ unsigned long resultSize);
// t = expr floordiv c <=> t = q, c * q <= expr <= c * q + c - 1
// A floordiv is thus flattened by introducing a new local variable q, and
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..bf7e3121ea32a 100644
--- a/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
+++ b/mlir/lib/Analysis/FlatLinearValueConstraints.cpp
@@ -36,11 +36,12 @@ using namespace presburger;
namespace {
// See comments for SimpleAffineExprFlattener.
-// An AffineExprFlattener extends a SimpleAffineExprFlattener by recording
-// constraint information associated with mod's, floordiv's, and ceildiv's
-// in FlatLinearConstraints 'localVarCst'.
+// An AffineExprFlattenerWithLocalVars extends a SimpleAffineExprFlattener by
+// recording constraint information associated with mod's, floordiv's, and
+// ceildiv's in FlatLinearConstraints 'localVarCst'.
struct AffineExprFlattener : public SimpleAffineExprFlattener {
-public:
+ using SimpleAffineExprFlattener::SimpleAffineExprFlattener;
+
// Constraints connecting newly introduced local variables (for mod's and
// div's) to existing (dimensional and symbolic) ones. These are always
// inequalities.
@@ -48,7 +49,7 @@ struct AffineExprFlattener : public SimpleAffineExprFlattener {
AffineExprFlattener(unsigned nDims, unsigned nSymbols)
: SimpleAffineExprFlattener(nDims, nSymbols),
- localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {}
+ localVarCst(PresburgerSpace::getSetSpace(nDims, nSymbols)) {};
private:
// Add a local variable (needed to flatten a mod, floordiv, ceildiv expr).
@@ -63,76 +64,144 @@ struct AffineExprFlattener : public SimpleAffineExprFlattener {
// Update localVarCst.
localVarCst.addLocalFloorDiv(dividend, divisor);
}
+
+ LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs,
+ ArrayRef<int64_t> rhs,
+ AffineExpr localExpr) override {
+ // AffineExprFlattener does not support semi-affine expressions.
+ return failure();
+ }
+};
+
+// A SemiAffineExprFlattener is an AffineExprFlattenerWithLocalVars that adds
+// conservative bounds for semi-affine expressions (given assumptions hold). If
+// the assumptions required to add the semi-affine bounds are found not to hold
+// the final constraints set will be empty/inconsistent. If the assumptions are
+// never contradicted the final bounds still only will be correct if the
+// assumptions hold.
+struct SemiAffineExprFlattener : public AffineExprFlattener {
+ using AffineExprFlattener::AffineExprFlattener;
+
+ LogicalResult addLocalIdSemiAffine(ArrayRef<int64_t> lhs,
+ ArrayRef<int64_t> rhs,
+ AffineExpr localExpr) override {
+ auto result =
+ SimpleAffineExprFlattener::addLocalIdSemiAffine(lhs, rhs, localExpr);
+ assert(succeeded(result) &&
+ "unexpected failure in SimpleAffineExprFlattener");
+ (void)result;
+
+ if (localExpr.getKind() == AffineExprKind::Mod) {
+ // Given two numbers a and b, division is defined as:
+ //
+ // a = bq + r
+ // 0 <= r < |b| (where |x| is the absolute value of x)
+ //
+ // q = a floordiv b
+ // r = a mod b
+
+ // Add a new local variable (r) to represent the mod.
+ unsigned rPos = localVarCst.appendVar(VarKind::Local);
+
+ // r >= 0 (Can ALWAYS be added)
+ localVarCst.addBound(BoundType::LB, rPos, 0);
+
+ // r < b (Can be added if b > 0, which we assume here)
+ ArrayRef<int64_t> b = rhs;
+ SmallVector<int64_t> bSubR(b);
+ bSubR.insert(bSubR.begin() + rPos, -1);
+ // Note: bSubR = b - r
+ // So this adds the bound b - r >= 1 (equivalent to r < b)
+ localVarCst.addBound(BoundType::LB, bSubR, 1);
+
+ // Note: The assumption of b > 0 is based on the affine expression docs,
+ // which state "RHS of mod is always a constant or a symbolic expression
+ // with a positive value." (see AffineExprKind in AffineExpr.h). If this
+ // assumption does not hold constraints (added above) are a contradiction.
+
+ return success();
+ }
+
+ // TODO: Support other semi-affine expressions.
+ return failure();
+ }
};
} // 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);
- // 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();
- }
+ auto flattenExprs = [&](AffineExprFlattener &flattener) -> LogicalResult {
+ // Use the same flattener to simplify each expression successively. This way
+ // local variables / expressions are shared.
+ for (auto expr : exprs) {
+ auto flattenResult = flattener.walkPostOrder(expr);
+ if (failed(flattenResult))
+ return failure();
+ }
- assert(flattener.operandExprStack.size() == exprs.size());
- flattenedExprs->clear();
- flattenedExprs->assign(flattener.operandExprStack.begin(),
- flattener.operandExprStack.end());
+ assert(flattener.operandExprStack.size() == exprs.size());
+ flattenedExprs->clear();
+ flattenedExprs->assign(flattener.operandExprStack.begin(),
+ flattener.operandExprStack.end());
- if (localVarCst)
- localVarCst->clearAndCopyFrom(flattener.localVarCst);
+ if (localVarCst)
+ localVarCst->clearAndCopyFrom(flattener.localVarCst);
- return success();
+ return success();
+ };
+
+ if (addConservativeSemiAffineBounds) {
+ SemiAffineExprFlattener flattener(numDims, numSymbols);
+ return flattenExprs(flattener);
+ }
+
+ AffineExprFlattener flattener(numDims, numSymbols);
+ return flattenExprs(flattener);
}
// 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 +710,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 +735,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 +751,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 +789,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
diff --git a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
index f8df34843a363..9c365376c84c9 100644
--- a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
+++ b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
@@ -7,9 +7,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
-
#include "mlir/Dialect/Vector/IR/VectorOps.h"
-
namespace mlir::vector {
FailureOr<ConstantOrScalableBound::BoundSize>
@@ -62,6 +60,11 @@ ScalableValueBoundsConstraintSet::computeScalableBound(
int64_t pos = scalableCstr.insert(value, dim, /*isSymbol=*/false);
scalableCstr.processWorklist();
+ // Check the resulting constraints set is valid.
+ if (scalableCstr.cstr.isEmpty()) {
+ return failure();
+ }
+
// 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) {
@@ -71,6 +74,11 @@ ScalableValueBoundsConstraintSet::computeScalableBound(
return p.first != scalableCstr.getVscaleValue() && !isStartingPoint;
};
scalableCstr.projectOut(projectOutFn);
+ // Also project out local variables (these are not tracked by the
+ // ValueBoundsConstraintSet).
+ for (unsigned i = 0, e = scalableCstr.cstr.getNumLocalVars(); i < e; ++i) {
+ scalableCstr.cstr.projectOut(scalableCstr.cstr.getNumDimAndSymbolVars());
+ }
assert(scalableCstr.cstr.getNumDimAndSymbolVars() ==
scalableCstr.positionToValueDim.size() &&
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index 94562d0f15a24..5f2016470b25f 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -1242,13 +1242,13 @@ LogicalResult SimpleAffineExprFlattener::visitMulExpr(AffineBinaryOpExpr expr) {
// variable in place of the product; the affine expression
// corresponding to the quantifier is added to `localExprs`.
if (!isa<AffineConstantExpr>(expr.getRHS())) {
+ SmallVector<int64_t, 8> mulLhs(lhs);
MLIRContext *context = expr.getContext();
AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols,
localExprs, context);
AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
localExprs, context);
- addLocalVariableSemiAffine(a * b, lhs, lhs.size());
- return success();
+ return addLocalVariableSemiAffine(mulLhs, rhs, a * b, lhs, lhs.size());
}
// Get the RHS constant.
@@ -1295,13 +1295,13 @@ LogicalResult SimpleAffineExprFlattener::visitModExpr(AffineBinaryOpExpr expr) {
// variable in place of the modulo value, and the affine expression
// corresponding to the quantifier is added to `localExprs`.
if (!isa<AffineConstantExpr>(expr.getRHS())) {
+ SmallVector<int64_t, 8> modLhs(lhs);
AffineExpr dividendExpr = getAffineExprFromFlatForm(
lhs, numDims, numSymbols, localExprs, context);
AffineExpr divisorExpr = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
localExprs, context);
AffineExpr modExpr = dividendExpr % divisorExpr;
- addLocalVariableSemiAffine(modExpr, lhs, lhs.size());
- return success();
+ return addLocalVariableSemiAffine(modLhs, rhs, modExpr, lhs, lhs.size());
}
int64_t rhsConst = rhs[getConstantIndex()];
@@ -1385,19 +1385,22 @@ SimpleAffineExprFlattener::visitConstantExpr(AffineConstantExpr expr) {
return success();
}
-void SimpleAffineExprFlattener::addLocalVariableSemiAffine(
- AffineExpr expr, SmallVectorImpl<int64_t> &result,
- unsigned long resultSize) {
+LogicalResult SimpleAffineExprFlattener::addLocalVariableSemiAffine(
+ ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr,
+ SmallVectorImpl<int64_t> &result, unsigned long resultSize) {
assert(result.size() == resultSize &&
"`result` vector passed is not of correct size");
int loc;
- if ((loc = findLocalId(expr)) == -1)
- addLocalIdSemiAffine(expr);
+ if ((loc = findLocalId(localExpr)) == -1) {
+ if (failed(addLocalIdSemiAffine(lhs, rhs, localExpr)))
+ return failure();
+ }
std::fill(result.begin(), result.end(), 0);
if (loc == -1)
result[getLocalVarStartIndex() + numLocals - 1] = 1;
else
result[getLocalVarStartIndex() + loc] = 1;
+ return success();
}
// t = expr floordiv c <=> t = q, c * q <= expr <= c * q + c - 1
@@ -1426,13 +1429,13 @@ LogicalResult SimpleAffineExprFlattener::visitDivExpr(AffineBinaryOpExpr expr,
// variable in place of the quotient, and the affine expression corresponding
// to the quantifier is added to `localExprs`.
if (!isa<AffineConstantExpr>(expr.getRHS())) {
+ SmallVector<int64_t, 8> divLhs(lhs);
AffineExpr a = getAffineExprFromFlatForm(lhs, numDims, numSymbols,
localExprs, context);
AffineExpr b = getAffineExprFromFlatForm(rhs, numDims, numSymbols,
localExprs, context);
AffineExpr divExpr = isCeil ? a.ceilDiv(b) : a.floorDiv(b);
- addLocalVariableSemiAffine(divExpr, lhs, lhs.size());
- return success();
+ return addLocalVariableSemiAffine(divLhs, rhs, divExpr, lhs, lhs.size());
}
// This is a pure affine expr; the RHS is a positive constant.
@@ -1503,11 +1506,14 @@ void SimpleAffineExprFlattener::addLocalFloorDivId(ArrayRef<int64_t> dividend,
// dividend and divisor are not used here; an override of this method uses it.
}
-void SimpleAffineExprFlattener::addLocalIdSemiAffine(AffineExpr localExpr) {
+LogicalResult SimpleAffineExprFlattener::addLocalIdSemiAffine(
+ ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs, AffineExpr localExpr) {
for (SmallVector<int64_t, 8> &subExpr : operandExprStack)
subExpr.insert(subExpr.begin() + getLocalVarStartIndex() + numLocals, 0);
localExprs.push_back(localExpr);
++numLocals;
+ // lhs and rhs are not used here; an override of this method uses them.
+ return success();
}
int SimpleAffineExprFlattener::findLocalId(AffineExpr localExpr) {
diff --git a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
index 87937591e60ad..6420c192b257d 100644
--- a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
+++ b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
@@ -151,8 +151,10 @@ ValueBoundsConstraintSet::Variable::Variable(AffineMap map,
[](Value v) { return Variable(v); })) {}
ValueBoundsConstraintSet::ValueBoundsConstraintSet(
- MLIRContext *ctx, StopConditionFn stopCondition)
- : builder(ctx), stopCondition(stopCondition) {
+ MLIRContext *ctx, StopConditionFn stopCondition,
+ bool addConservativeSemiAffineBounds)
+ : builder(ctx), stopCondition(stopCondition),
+ addConservativeSemiAffineBounds(addConservativeSemiAffineBounds) {
assert(stopCondition && "expected non-null stop condition");
}
@@ -174,11 +176,19 @@ static void assertValidValueDim(Value value, std::optional<int64_t> dim) {
void ValueBoundsConstraintSet::addBound(BoundType type, int64_t pos,
AffineExpr expr) {
+ // Note: If `addConservativeSemiAffineBounds` is true then the bound
+ // computation function needs to handle the case that the constraints set
+ // could become empty. This is because the conservative bounds add assumptions
+ // (e.g. for `mod` it assumes `rhs > 0`). If these constraints are later found
+ // not to hold, then the bound is invalid.
LogicalResult status = cstr.addBound(
type, pos,
- AffineMap::get(cstr.getNumDimVars(), cstr.getNumSymbolVars(), expr));
+ AffineMap::get(cstr.getNumDimVars(), cstr.getNumSymbolVars(), expr),
+ addConservativeSemiAffineBounds
+ ? FlatLinearConstraints::AddConservativeSemiAffineBounds::Yes
+ : FlatLinearConstraints::AddConservativeSemiAffineBounds::No);
if (failed(status)) {
- // Non-pure (e.g., semi-affine) expressions are not yet supported by
+ // Not all semi-affine expressions are not yet supported by
// FlatLinearConstraints. However, we can just ignore such failures here.
// Even without this bound, there may be enough information in the
// constraint system to compute the requested bound. In case this bound is
diff --git a/mlir/test/Dialect/Vector/test-scalable-bounds.mlir b/mlir/test/Dialect/Vector/test-scalable-bounds.mlir
index d549c5bd1c378..673e03f05c1b8 100644
--- a/mlir/test/Dialect/Vector/test-scalable-bounds.mlir
+++ b/mlir/test/Dialect/Vector/test-scalable-bounds.mlir
@@ -159,3 +159,59 @@ func.func @non_scalable_code() {
}
return
}
+
+// -----
+
+#remainder_start_index = affine_map<()[s0] -> (-(1000 mod s0) + 1000)>
+#remaining_iterations = affine_map<(d0) -> (-d0 + 1000)>
+
+// CHECK: #[[$REMAINDER_START_MAP:.*]] = affine_map<()[s0] -> (-(1000 mod s0) + 1000)>
+// CHECK: #[[$SCALABLE_BOUND_MAP_4:.*]] = affine_map<()[s0] -> (s0 * 8 - 1)>
+
+// CHECK-LABEL: @test_scalable_remainder_loop
+// CHECK: %[[VSCALE:.*]] = vector.vscale
+// CHECK: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_4]]()[%[[VSCALE]]]
+// CHECK: "test.some_use"(%[[SCALABLE_BOUND]]) : (index) -> ()
+func.func @test_scalable_remainder_loop() {
+ %c8 = arith.constant 8 : index
+ %c1000 = arith.constant 1000 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %vscale, %c8 : index
+ %0 = affine.apply #remainder_start_index()[%c8_vscale]
+ scf.for %arg1 = %0 to %c1000 step %c8_vscale {
+ %remaining_iterations = affine.apply #remaining_iterations(%arg1)
+ // 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, type = "UB", vscale_min = 1 : i64, vscale_max = 16 : i64}> : (index) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ }
+ return
+}
+
+// -----
+
+#unsupported_semi_affine = affine_map<()[s0] -> (s0 * s0)>
+
+func.func @unsupported_semi_affine() {
+ %vscale = vector.vscale
+ %0 = affine.apply #unsupported_semi_affine()[%vscale]
+ // expected-error @below{{could not reify bound}}
+ %bound = "test.reify_bound"(%0) <{scalable, type = "UB", vscale_min = 1 : i64, vscale_max = 16 : i64}> : (index) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ return
+}
+
+// -----
+
+#map_mod = affine_map<()[s0] -> (1000 mod s0)>
+
+func.func @unsupported_negative_mod() {
+ %c_minus_1 = arith.constant -1 : index
+ %vscale = vector.vscale
+ %negative_vscale = arith.muli %vscale, %c_minus_1 : index
+ %0 = affine.apply #map_mod()[%negative_vscale]
+ // expected-error @below{{could not reify bound}}
+ %bound = "test.reify_bound"(%0) <{scalable, type = "UB", vscale_min = 1 : i64, vscale_max = 16 : i64}> : (index) -> index
+ "test.some_use"(%bound) : (index) -> ()
+ return
+}
More information about the Mlir-commits
mailing list