This is an archive of the discontinued LLVM Phabricator instance.

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

Authored by springerm on Mar 9 2023, 3:39 AM.

Details

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

Diff Detail

Event Timeline

springerm created this revision.Mar 9 2023, 3:39 AM
springerm requested review of this revision.Mar 9 2023, 3:39 AM
Herald added a project: Restricted Project. · View Herald TranscriptMar 9 2023, 3:39 AM
springerm added inline comments.Mar 9 2023, 3:49 AM
mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp
131 ↗(On Diff #503721)

I have not found a good way yet to do this in the transform dialect. This change does not add a transform. It just adds an interface and its implementation. I think there are certain things that are easier (less code in the test file) to test via transform ops and other things are better to test via test passes.

I could add a transform.structured.reify_bounds op that takes a value handle and produces a value handle. This, combined with another new transform dialect op for "replaceAllUsesWith" could be a replacement for what I'm doing in TestLinalgTransforms.cpp. There would be a lot of boilerplate code though.

If this has to be a transform dialect op, ideally we would add an op that is also useful outside of tests. The main question is how are we going to use this functionality from the transform dialect. Do we have a use case for things such as linalg::computeUpperBoundForIndex as a transform op or this usually part of another transform op implementation? In the latter case, it may be better to keep this test in TestLinalgTransforms.cpp.

springerm retitled this revision from [mlir][linalg] Add ValueBoundsOpInterface to [mlir][linalg] Add ValueBoundsOpInterface and tensor dialect op impl.Mar 9 2023, 8:16 AM

I took a first look and this looks really promising! Thank you so much for working on this so quickly.
Before we move forward with deeper reviews and comments, do you plan to send and RFC? I think this interface may benefit from one as it will impact multiple dialects. It would recommend so even if it's just an informative one.

Thanks!
Diego

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h
24 ↗(On Diff #503721)

Perhaps adding more context here about terms used in the public APIs would help (e.g., worklist, used value)

27 ↗(On Diff #503721)

Don't have a mental model yet but this comment makes me think that this class holds a state (mapping) that may become invalid if the IR changes. Is that the case?

35 ↗(On Diff #503721)

What about having to methods, one for values and one for dims? Clearer, less error-prone for the user and we wouldn't have to expose kIndexValue publicly.
Another option is using an std::optional. I'm not very excited about using a -1 but feel free to keep it if you think it simplifies the implementation/other APIs.

74 ↗(On Diff #503721)

(Not sure I understand yet what a used value is)
It looks like this getExpr is returning an AffineExpr but also doing something else underneath. Make the name more explicit?

117 ↗(On Diff #503721)

Would it make sense to build your own builder? If this is kind of an analysis, it may happen that the user doesn't have a builder. We should only need an MLIRContext.

mlir/test/Dialect/Linalg/value-bounds-reification.mlir
1 ↗(On Diff #503721)

typo

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp
228 ↗(On Diff #503721)

Doc about reifyToFuncArgs?
Doc about what this function does? I expected this to just print the bound but that's not the case.

springerm marked 2 inline comments as done.Mar 12 2023, 6:00 AM

Before we move forward with deeper reviews and comments, do you plan to send and RFC? I think this interface may benefit from one as it will impact multiple dialects. It would recommend so even if it's just an informative one.

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

mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h
27 ↗(On Diff #503721)

Yes, that's correct. Also note that the only constructor of ValueBoundsConstraintSet is private and reifyBound does not return a ValueBoundsConstraintSet. I.e., there's no way for the user to get hold of a ValueBoundsConstraintSet object.

74 ↗(On Diff #503721)

The argument name val may be confusing. This function simply returns an AffineConstExpr. Users could build one themselves without using this API, but then they would have to create a Builder. This is just a small convenience function.

117 ↗(On Diff #503721)

This builder is created (from the MLIRContext object inside the ValueDim that is passed to the constructor) and owned by the ValueBoundsConstraintSet. It is just stored here so that we don't create a new builder every time an AffineExpr is created. Users cannot pass a builder to this class. (But they have to pass one to reifyBounds.)

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp
228 ↗(On Diff #503721)

Doc is in the Option definition in line 134: If set to true, values will be reified in terms of FuncOp bbArgs. (Otherwise, in terms of any values that are not owned by the same op; i.e., similar to reifyResultShapes.) Will add more docs to the reifyValueBounds function.

mehdi_amini added inline comments.Mar 12 2023, 6:07 AM
mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h
23 ↗(On Diff #503721)

I don't know what "columns" refers to? I suspect some examples here on top of what Diego mentions below would help

27 ↗(On Diff #503721)

I was about to ask the same thing: this is a map from SSA value to constraints, and it's not clear to me who holds this set and how to manage invalidation / recalculation?

44 ↗(On Diff #503721)
mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.td
14 ↗(On Diff #503721)

Why is this inside Linalg? Is it because of "reifyBound" depending on tensor/memref dialects? (see my comment there)

This is an major red flag to me right now.

30 ↗(On Diff #503721)

It's not clear to me what the "value" is here, is this a result of the operation? The doc does not seem to say.

In this case can we pass an index instead? Or an OpResult?

Edit: looking at the implementation, it can also be the parent op in case this is a block argument! I suspect the current implementation does not seem suitable for a CFG?

mlir/include/mlir/InitAllDialects.h
136

There were supposed to be a mechanism to make this all safe that was discussed when you started using this pattern initially, was it implemented?

mlir/lib/Dialect/Linalg/Transforms/ValueBoundsOpInterface.cpp
303 ↗(On Diff #503721)

Can we replace this with an interface method on ShapedTypeInterface and avoid hard-depending on specific dialect implementations?

mlir/test/Dialect/Linalg/value-bounds-op-interface-impl.mlir
2 ↗(On Diff #503721)

Please don't use -allow-unregistered-dialect anywhere, we should never use this in the codebase.

springerm marked 2 inline comments as done.Mar 12 2023, 6:32 AM
springerm added inline comments.
mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.td
14 ↗(On Diff #503721)

No particular reason. It should not be in LinalgTransforms, I just had to put it somewhere for the moment.

My preferred solution would be:

  • Use IntegerRelation instead of FlatAffineValueConstraints. (IntegerRelation/IntegerPolyhedron needs a few functions that are currently defined on FlatAffineValueConstraints, but I think this is doable.) ValueBoundsOpInterface and ValueBoundsConstraintSet can then be in mlir/Interfaces.
  • ValueBoundsConstraintSet::reifyBound becomes a standalone static function in AffineTransforms (or AffineDialect if we can find a way to create memref.dim/tensor.dim without depending on TensorDialect/MemRefDialect).
  • We could add a second reifyBound function to ArithTransforms (that will create std.addi etc. instead of affine.apply).
30 ↗(On Diff #503721)

In most cases it would be an OpResult. There are a few cases where we would want to support a BlockArgument. E.g., it could be the IV of an scf.for. I think in case of a BlockArgument, this should be restricted to "entry" block arguments of regions.

(I think the current worklist approach is too simple to handle cyclic CFGs etc.; it could be extended in the future if there are use cases. Same goes for iter_args of scf.for etc.: D145804 adds an interface impl for scf.for, but iter_args are not supported.)

I will update the docs accordingly.

dcaballe added inline comments.Mar 12 2023, 11:32 PM
mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h
27 ↗(On Diff #503721)

Same question as Mehdi.

This looks like an analysis to me. I haven't seen analysis passes implemented in MLIR but we have the infra: https://mlir.llvm.org/docs/PassManagement/#analysis-management. I wonder if this would be a good fit for an analysis.

117 ↗(On Diff #503721)

(But they have to pass one to reifyBounds.)

Exactly. I had to create a dummy builder just to invoke reifyBounds.

mehdi_amini added inline comments.Mar 14 2023, 8:23 AM
mlir/include/mlir/InitAllDialects.h
136

Found the revision: https://reviews.llvm.org/D120368

I would really like to see this getting through before we add more "external interface" in an unsafe fashion. Can you help with this?

springerm edited the summary of this revision. (Show Details)Mar 16 2023, 8:00 AM
springerm marked 7 inline comments as done.
springerm added inline comments.
mlir/include/mlir/Dialect/Linalg/Transforms/ValueBoundsOpInterface.h
35 ↗(On Diff #503721)

std::optional did not work the DenseMap, but we can at least use it for the public API.

springerm added inline comments.Mar 16 2023, 8:02 AM
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)
springerm updated this revision to Diff 505809.Mar 16 2023, 8:02 AM
springerm marked 2 inline comments as done.
springerm edited the summary of this revision. (Show Details)

address some comments and move to mlir/Interfaces

springerm retitled this revision from [mlir][linalg] Add ValueBoundsOpInterface and tensor dialect op impl to [mlir][Interfaces] Add ValueBoundsOpInterface and tensor dialect op impl.Mar 16 2023, 8:03 AM
springerm edited the summary of this revision. (Show Details)Mar 16 2023, 8:04 AM
springerm updated this revision to Diff 506110.Mar 17 2023, 9:24 AM

Remove loc

springerm added inline comments.Mar 20 2023, 4:28 AM
mlir/include/mlir/InitAllDialects.h
136

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

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

mlir/include/mlir/InitAllDialects.h
136

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

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

jpienaar added inline comments.Mar 29 2023, 9:37 PM
mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h
52

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?

59

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?

64

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
springerm marked 12 inline comments as done.Mar 30 2023, 9:22 AM
springerm added inline comments.
mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h
52

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

I added an assertion in ValueBoundsConstraintSet::computeBound.

59

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 updated this revision to Diff 509699.Mar 30 2023, 9:22 AM
springerm marked 8 inline comments as done.

address most comments

springerm added inline comments.Mar 31 2023, 3:33 AM
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.

ubfx added a subscriber: ubfx.Apr 3 2023, 10:40 PM
jpienaar accepted this revision.Apr 5 2023, 2:58 PM

Nice, I like the new API too.

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

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 ?

This revision is now accepted and ready to land.Apr 5 2023, 2:58 PM
dcaballe accepted this revision.Apr 5 2023, 5:43 PM
This revision was landed with ongoing or failed builds.Apr 5 2023, 5:57 PM
This revision was automatically updated to reflect the committed changes.
springerm marked 5 inline comments as done.

Thanks for the thorough review!

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

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().