[mlir][Interfaces] Add ValueBoundsOpInterface and tensor dialect op impl

[springerm]

Summary

Ops can implement this interface to specify lower/upper bounds for their result values and block arguments. Bounds can be specified for:

• Index-type values
• Dimension sizes of shapes values

The bounds are added to a constraint set. Users can query this constraint set to compute bounds wrt. to a user-specified set of values. Only EQ bounds are supported at the moment.

A helper function is added to the Affine dialect to reify a computed bounds with an affine.apply op.

This revision also contains interface implementations for various tensor dialect ops, which illustrates how to implement this interface.

RFC: https://discourse.llvm.org/t/rfc-valueboundsopinterface-compute-bounds-of-index-values-dyn-dim-sizes/69174

Depends On: D146869

[springerm]

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h

117 ↗

(On Diff #503721)

That is because reifyBounds can be used to reify non-constant bounds. I.e., the bound may be the result of an affine.apply op that computes the bound. You're using it to compute constant bounds (in which case no builder would be needed), but reifyBounds is more general.

I'm going to add a new option/overload for computing an explicitly constant bound (without a builder) in a subsequent revision.

mlir/lib/Dialect/Linalg/Transforms/ValueBoundsOpInterface.cpp

303 ↗

(On Diff #503721)

D146052

[springerm]

address some comments and move to mlir/Interfaces

[springerm]

update

[springerm]

rebase

[springerm]

Remove loc

[springerm]

mlir/include/mlir/InitAllDialects.h

138

D146414

[springerm]

rebase

[dcaballe]

LGTM. All my comments have been resolved. Please, wait for Mehdi to provide feedback on the remaining open questions/secondary PRs.

[mehdi_amini]

I'm curious about @jpienaar 's thoughts on the overall direction here

mlir/include/mlir/InitAllDialects.h

138

I am not sure what you have in mind for D146414, I meant helping reviewing D120368 to get it to land.

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

460 ↗

(On Diff #506945)

Can you commit this separately? It's not clear to me that it relates to the current patch

[springerm]

rebase

[jpienaar]

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

53

So value and dim may be provided as long as value is not index type? E.g., if value is an i32, then you can have both?

60

It is difficult to know what this function will do and stopCondition represents - it feels like base case of recursion, but its unclear here. Could you perhaps expand this to provide more information? How many stop conditions are actually here?

65

And ValueBoundsConstraintSet is not visible here so that one could use the stop condition type defined there?

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

37

That's just being conservative right? Could this be retained in more cases?

45

OOC are all BoundType's allowed?

62

I almost wonder if spelling out boundtype in name isn't more intuitive, e.g.,

addUpperBound, addLowerBound, addEquality (last one sounds less good than other two).

Ha, makes me wonder about operators boundset.bound(value) <= dim / boundset.bound(value)[dim] <= expr - which does also translate closer to

addBound(value, UB, expr) - so it reads add bound for value of upper bound with expr.

Seeing presburger in the API just feels weird to me.

80

These aren't bounds right? This feels natural with equality bounds/solution of constraint set.

92

Can we choose any other sentinel (-56) or differentiate with type (optional)? (-1 has been used in sentinel in too many places ...)

102

Mehdi commented on this as well, but this is first mention of column, might be worthwhile to introduce here.

148

Could this be along with destination style ops or in another header?

mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp

24

Could you expand here? Its not quite intuitive what the effect is (well and presents a nicely documented usage of the interface).

54

New form is

cast<Foo>(...) / isa<Foo>(...)

(https://mlir.llvm.org/deprecation/)

58

(your other change will improve this, I also toyed around with having a shaped.dim op that canonicalizes on input type without dependency on dialect https://gist.github.com/jpienaar/82c79ae73b2f46be4c9f365334bf8dc3)

mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp

20

I'd really prefer if we did this in a more declarative manner so that this could be generated and we could even have a dynamic form of this (similar to how torch-mlir does shape functions in spirit). But more work is needed for that and we aren't there yet, and that can't be a blocker here.

74

This feels like exactly what we'd represent in the shape (and type inference) function.

120

Why declare this as an external interface inside the dialect directory? Are alternative implementations expected? This going to be optionally added?

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

65

Could we capture the error in op and surface it when queried or some such? Or is it assumed that verification would have caught all instances where bounds could fail?

150

https://llvm.org/docs/CodingStandards.html#don-t-evaluate-end-every-time-through-a-loop

[springerm]

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

53

i32 is not supported (added with D146870). Only index-type values and shaped values are supported.

I added an assertion in ValueBoundsConstraintSet::computeBound.

60

Added an example. Some of the next revisions also add overloads, so that users can use this interface without specifying a stopCondition.

• D146296 makes the stopCondition optional for constant bounds.
• D146306 adds an overload for a given list of dependencies.
• D143910 adds an overload for a given list of independencies.

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

37

Correct. Adding new IR is OK. Expanded the comment.

45

Only EQ at the moment. There is an assertion in the implementation that LB/UB are not yet implemented. The other bound types are added in D145787 to keep this revision a bit smaller.

62

This is a nice API. And it allows us to specify inclusive/exclusive bounds via </<= etc.

I updated just this revision for now, will update the dependent ones later.

80

The getExpr function can be used to provide an expression that can be passed to addBound. (The overload that take an AffineExpr.)

92

I wanted to use std::optional but for some reason, an std::optional cannot be used as the key for a DenseMap:

DenseMap.h:469:22: error: no member named 'getEmptyKey' in 'llvm::DenseMapInfo<std::optional<long>>'

148

This must be in some place where both DestinationStyleOpInterface and ValueBoundsOpInterface are available. (Kind of like cross-dialect dependencies, but this one is cross-interfaces.)

ValueBoundsOpInterface seems a bit better than putting it in DestinationStyleOpInterface.h because when DestinationStyleOpInterface is a more "fundamental" interface.

mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp

20

What do you mean by "more declarative"?

120

It could be implemented directly on the op. Alternative interface implementations are not possible in MLIR (not even with external models).

The reason why I put it here is separation of concerns. So that TensorOps.cpp/TensorOps.td is not getting any longer. This keeps TensorOps.td clean of any "optional/non-fundamental" interfaces and easy to understand. I think certain interfaces are very important and should be implemented on the op directly (e.g., OffsetSizeAndStrideOpInterface), other just clutter the file (like this one).

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

65

This is just a sanity check. This should always return success. If it fails, it indicates that there is a bug in FlatLinearConstraints or the ValueBoundsOpInterface implementation.

[springerm]

address most comments

[springerm]

update

[springerm]

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

92

I updated a few more places in the internal API. kIndexValue is now just used when interacting with the DenseMap.

[jpienaar]

Nice, I like the new API too.

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

53

I meant something else: is this optional simulating a sum type? E.g., is only one of value or dim allowed to be set and an optional is being used? So a ValueOrDim bound is actually being passed in using more raw types (well a std::variant<Value, int64_t> is probably not much more higher level than std::optional, so raw perhaps not best word)

mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp

20

Something that could allow us to capture the information in one way but use both places rather than implement twice. I feel like if we had it such that you could run an abstract interpretation over the shape inference functions, you could have gotten some of these for free (and vice versa). [I'll probably prod at this in a month or so, so no blocker here, just me wishing that was done already :)]

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

50

Following new style preference/deprecation, dyn_cast<ShapedType>(value.getType()) here and above.

mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp

30

Update description ?

[springerm]

Thanks for the thorough review!

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

53

std::variant is like a union right? This is different here. This function can be called in two ways:

• value has index type and dim = nullopt
• value has a shaped type and dim != nullopt

So value is always needed and whether dim should be specified or not depends on value.getType().