Status update: I uploaded a revision that addresses some of the comments mentioned in here. ValueBoundsOpInterface no longer builds IR. It is now analysis only. There are two helpers (in the respective dialects) that manifest a computed bound in IR; either as an affine.apply op or as a combination of arith dialect ops.
tensorflow/tensorflow/compiler/mlir/tools/kernel_gen/transforms/same_shape_propagation.cc
This component builds equivalence classes of SSA values that have the same shape.
ValueBoundsOpInterface could be a complete upstream replacement for this component. Some issues with same_shape_propagation.cc that could be addressed with ValueBoundsOpInterface:
- Hard-coded list of supported ops (not interface-based).
- Can only handle cases where the shape of a result is the same as the shape of another value. No arithmetics possible. E.g.,
tensor.padcould not be handled bysame_shape_propagation.cc.
There’s also this comment. ValueBoundsOpInterface would address that comment.
// A basic shape equality inference. This should be superceeded by a proper
// inference once available. Until then, we just build this out to the needs of
// the kernel generator project.
xla/xla/mlir_hlo/mhlo/analysis/shape_component_analysis.cc
This is a symbolic interpreter for MHLO ops. It can compute dimension sizes (GetShapeInfo) and values of SSA values (GetValueInfo).
The handling of SSA values is slightly different from ValueBoundsOpInterface: it can be used to compute the elements stored in a shape tensor, whereas ValueBoundsOpInterface only works on index/integer-typed SSA values at the moment.
There are many similarities in the design: both approaches use a worklist-driven approach to traverse the backward slice. Both utilize AffineExpr to encode bounds. The SymbolicExpr data structure is very similar to the combination of AffineMap and ValueDimList in ValueBoundsOpInterface.
shape_component_analysis.cc is currently not interface-based and has a hard-coded list of supported ops. If we add support for shape tensor values to ValueBoundsOpInterface, it could be an upstream replacement for shape_component_analysis.cc.
MLIR InferIntRangeInterface
This interface can be used to constant integer LB and UB for integer-typed SSA values. ValueBoundsOpInterface is a more generalized variant, which also compute non-constant bounds wrt. SSA values.
The data structures are different. InferIntRangeInterface uses a lattice. I’m not super familiar with this data structure. ValueBoundsOpInterface uses a constraint set (FlatLinearConstraints).
ValueBoundsOpInterface can be seen as a generalization of InferIntRangeInterface. I think we should try to merge the two. I think ValueBoundsConstraintSet::computeConstantBound is doing roughly the same as querying a single bound from the lattice. Maybe the authors of InferIntRangeInterface and related infrastructure could also weigh in a bit. @krzysz00 @Mogball
InferIntRangeInterface is aware of integer type ranges:
Bounds are tracked separately for the signed and unsigned
interpretations of a value, which ensures that the impact of
arithmetic overflows is correctly tracked during the analysis.
ValueBoundsOpInterface does currently not take this into account. The underlying constraint set is based on MPInt (similar to APInt). Once a bound is added to the constraint set, the type information is “erased”; there are no restrictions about how small/large a value can get.
In general, when analyzing a non-constant value, it may be impossible to know if a value overflows at runtime or not. For the special case where a constant bound is computed, it should be possible to take into account such overflows. But this will require some extra work.
I have a preliminary revision that adds support for branches to ValueBoundsOpInterface (arith.select, scf.if). These are tricky because their value bounds cannot be expressed as linear constraints. E.g., for %0 = arith.select %cond, %a, %b : index, the lower bound of constraint would be %0 >= min(lb(%a), lb(%b)). For such ops, ValueBoundsOpInterface can infer constant bounds (if there are any), so the constraint set-based approach can still do as well as the lattice-based approach (but no better).
MLIR ReifyRankedShapedTypeOpInterface
This is a special case of ValueBoundsOpInterface, where a bound is computed and reified as IR in terms of the op’s operands. The ReifyRankedShapedTypeOpInterface would no longer be needed. mlir::reifyResultShapes can directly query the ValueBoundsOpInterface. This would also allow dialects to implement the interface without depending on the arith/affine dialect.
There is a related discussion about inefficiencies when using ReifyRankedShapedTypeOpInterface to reify the shape of a value over a chain of multiple ops, which would be improved with ValueBoundsOpInterface: Is there a way to avoid creating instructions that are going to be CSE-ed anyway