Summary
Provide an op interface (ValueBoundsOpInterface) that adds constraints about an OpResult / BlockArgument into a constraint set (FlatAffineValueConstraints). Constraints can be added for index-typed values and dimension sizes of shaped values. The constraint set can be queried to compute lower/upper/equality bounds of index-typed values or dynamic dimension sizes in terms of a specified set of values.
Prototype: ⚙ D145681 [mlir][linalg] Add ValueBoundsOpInterface and tensor dialect op impl. Stack of revisions contains interface implementations for various ops.
Example use cases: Padding dynamic tensor/memref dimensions to static sizes (to make IR vectorizable, etc.), hoisting of buffers with dynamic sizes from loops, vector masking. Some concrete example test case are here.
There have been attempts to implement similar functionality in MLIR and downstream projects (ONNX, TensorFlow, IREE). ValueBoundsOpInterface could also generalize and replace the current ReifyRankedShapedTypeOpInterface.
Interface Methods
void populateBoundsForIndexValue(Value, ValueBoundsConstraintSet &)void populateBoundsForShapedValueDim(Value, int64_t dim, ValueBoundsConstraintSet &)
ValueBoundsConstraintSet is a thin wrapper around FlatAffineValueConstraints. It maintains a mapping of columns in constraint set to constrained values (index-typed SSA values or dims of shaped values). The mapping of columns to SSA values in FlatAffineValueConstraint cannot be used because it does not support “dim of shaped value”.
ValueBoundsConstraintSet also provides helper functions for retrieving constrained values/dims (getExpr) from the constraint set as AffineExprs and adding new bounds to the constraint set (addBound). In particular, AffineExprs make it easy to implement this interface with few lines of code.
ValueBoundsConstraintSet maintains an internal worklist for SSA values that were used in bounds of other values but who’s value bounds are not yet known. The class provides a static helper function reifyBound that iteratively processes values in the backward slice (reverse use-def chain) of a given value until a stop condition is met. It then projects out all columns for which the stop condition is not met and computes a bound as an IntegerAttr or Value. In the latter case, this can be an already existing SSA value (no new ops created) or the result of an affine.apply operation.
/// Reify a bound for the given index-typed value or shape dimension size in
/// terms of SSA values for which `stopCondition` is met.
static FailureOr<OpFoldResult>
ValueBoundsConstraintSet::reifyBound(
OpBuilder &b, Location loc, presburger::IntegerPolyhedron::BoundType type,
Value value, int64_t dim, function_ref<bool(Value)> stopCondition);
Example: tensor.pad
Compare this with the ReifyRankedShapedTypeOpInterface implementation of tensor.pad, which is doing roughly the same thing.
struct PadOpInterface
: public ValueBoundsOpInterface::ExternalModel<PadOpInterface, PadOp> {
void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
ValueBoundsConstraintSet &cstr) const {
auto padOp = cast<PadOp>(op);
assert(value == padOp.getResult() && "invalid value");
AffineExpr expr = cstr.getExpr(padOp.getSource(), dim) +
cstr.getExpr(padOp.getMixedLowPad()[dim]) +
cstr.getExpr(padOp.getMixedHighPad()[dim]);
cstr.addBound(IntegerPolyhedron::BoundType::EQ, value, dim, expr);
}
};
Example: scf.for
void populateBoundsForIndexValue(Operation *op, Value value,
ValueBoundsConstraintSet &cstr) const {
auto forOp = cast<ForOp>(op);
// Only IV is supported at the moment.
if (value != forOp.getInductionVar())
return;
// TODO: Take into account step size.
cstr.addBound(IntegerPolyhedron::BoundType::LB, value,
cstr.getExpr(forOp.getLowerBound()));
cstr.addBound(IntegerPolyhedron::BoundType::UB, value,
cstr.getExpr(forOp.getUpperBound()));
}
Example: API usage
Reify an upper bound for the size of dimension 0 of tensorValue in terms of vals or return failure if no such bound could be computed.
SmallVector<Value> vals = ...;
auto stopCondition = [&](Value v) { return llvm::contains(vals, v); };
FailureOr<OpFoldResult> bound = ValueBoundsConstraintSet::reifyBound(
b, loc, BoundType::UB, tensorValue, /*dim=*/0, stopCondition);
Generalization of ReifyRankedShapedTypeOpInterface
When reifying the dynamic shape of opResult, ValueBoundsOpInterface::reifyBound can be used as a replacement for the current reifyResultShapes when using the following stop condition.
auto stopCondition = [&](Value v) {
// Reify in terms of SSA values that are different from `value`.
return v != opResult;
};
Pros:
- Reifies only the specified dimension of the specified value, not every dimension of every result.
- The
ValueBoundsOpInterfaceimplementation is typically shorter (LoC) than theReifyRankedShapedTypeOpInterfaceimplement. E.g., compare the abovetensor.padimplementation with the corresponding implementation of the existing interface. - No code/functionality duplication between
ValueBoundsOpInterfaceandReifyRankedShapedTypeOpInterface. - If
ValueBoundsOpInterfaceis located inmlir/Transforms, ops can implement the interface directly, without having to depend onAffineDialectetc. (We would not longer need a separate build target inDialect/Tensoretc.)
Cons:
-
ValueBoundsOpInterfaceimplementations do not build IR. This is done as part ofValueBoundsOpInterface::reifyBound. This helper function must depend on at least one dialect (affineorarith) to materialize the IR that computes the shape. Therefore, it cannot be part ofmlir/Interfaces. (If we useIntegerPolyhedroninstead ofFlatAffineValueConstraints, the interface itself can, just thereifyBoundcannot.reifyBoundcould be placed inAffineTransformsand/orArithTransforms.) - Creating a constraint set adds some overhead. If this overhead is too large, we could add a “fast path” specifically for
reifyResultShapes, where we do not create a constraint set. (The constraint set is not really needed because no columns are projected out.)