Page MenuHomePhabricator

[MLIR] Introduce join/meet type interface.
Needs ReviewPublic

Authored by arames on May 4 2021, 2:25 PM.

Details

Summary

Introduce ::mlir::join(Type, Type) and ::mlir::meet(Type, Type), for the
partial order "less specialized than or equal".
Types can participate by implementing the JoinMeetTypeInterface. Today,
tensor types and memref types participate.

Diff Detail

Event Timeline

There are a very large number of changes, so older changes are hidden. Show Older Changes
silvas added inline comments.May 4 2021, 4:22 PM
mlir/test/Interfaces/JoinMeetTypeInterface/TestJoinMeetTypeInterface.cpp
71 ↗(On Diff #342867)

mem8 is too terse. memspace8?

silvas added inline comments.May 4 2021, 4:27 PM
mlir/lib/IR/BuiltinTypes.cpp
516

For RankedTensorType, you also need to check the "encoding"

rriddle added inline comments.May 4 2021, 4:34 PM
mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
55

Can you rename these to something specific to types? meet is a very general name that I wouldn't want polluting the mlir namespace. (Same with join)

mlir/test/Interfaces/JoinMeetTypeInterface/TestJoinMeetTypeInterface.cpp
23 ↗(On Diff #342867)

This file is a bit of an awkward thing to look at. Can you just write a simple pass that checks for a unregistered op with two types and tries to join or meet them? You can then have actual IR instead of parsing it yourself on the fly and using macros (which has a bad code smell to it).

Also, if you were to keep this kind of test it would make more sense as a unit test than a lit test IMO.

vinograd47 added inline comments.
mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
31

is non-ANSI character. AFAIK, they might cause issues with some compilers.

arames marked 14 inline comments as done.May 5 2021, 3:31 PM

Thanks for the review.

mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
55

Renamed to joinTypes and meetTypes.

mlir/include/mlir/Interfaces/JoinMeetTypeInterface.td
34

Since I already say "the contained type is null", I changed to say

[...], and someType and otherType do not join; that is, there is no type that is both less specialized than someType and less specialized than otherType.

mlir/lib/IR/BuiltinTypes.cpp
516

I missed that this was available on RankedTensorType.
Since it is not on UnrankedTensorType, I assume we can safely join tensor<*xf32> and tensor<2xf32, #any_encoding>, correct ?

mlir/test/Interfaces/JoinMeetTypeInterface/TestJoinMeetTypeInterface.cpp
23 ↗(On Diff #342867)

Moved to a pass as you suggested.

71 ↗(On Diff #342867)

Not applicable now that this file was removed in favor of an MLIR filecheck test.

silvas added inline comments.May 5 2021, 3:35 PM
mlir/lib/IR/BuiltinTypes.cpp
516

If there is a non-null encoding then joining with UnrankedTensorType should return null (the "encoding" is considered load-bearing for correctness and must be preserved).

silvas added a comment.May 5 2021, 3:36 PM

btw, I think you might need to reupload the patch. I don't see the new code.

arames marked 5 inline comments as done.May 5 2021, 3:37 PM

btw, I think you might need to reupload the patch. I don't see the new code.

Double-checking before uploading, sorry about that. I realize it's probably better to publish comments only after uploading the update in phabricator.

arames added inline comments.May 5 2021, 4:09 PM
mlir/lib/IR/BuiltinTypes.cpp
516

Conversely, I suppose meetTypes(tensor<*xi32>, tensor<?xi32, #some_map>) should return tensor<?xi32, #some_map> then.

This leads to some funny asymmetry.

join(tensor<1xi32, #encoding1>, tensor<2x2xi32>) == tensor<*xi32>
join(tensor<1xi32, #encoding1>, tensor<2x2xi32, #encoding1>) == <null> != tensor<*xi32, #encoding1>
arames added inline comments.May 5 2021, 4:12 PM
mlir/lib/IR/BuiltinTypes.cpp
516

Or maybe this is simply not possible due to how encoding works. Looking into it.

silvas added inline comments.May 5 2021, 4:19 PM
mlir/lib/IR/BuiltinTypes.cpp
516

It should be: join(tensor<1xi32, #encoding1>, tensor<2x2xi32>) == null

You can consider the lattice as a product lattice of tuples (TensorType's ignoring encoding, encoding), where encoding for the moment is a constant lattice (either top, bottom, or a single element set).

That is, unless two types have the same non-null encoding, their join/meet is null. (a null encoding is its own distinct encoding, not "any encoding")

mehdi_amini added inline comments.May 5 2021, 4:51 PM
mlir/lib/IR/BuiltinTypes.cpp
516

Having a shape-centric interface be bothered by encodings seems like a red flag to me.

It seems to me that this indicates that working with "Type" isn't the right level of abstraction and the API should work with a "shape" object instead.

silvas added inline comments.May 5 2021, 4:53 PM
mlir/lib/IR/BuiltinTypes.cpp
516

There's nothing shape-centric about this interface.

arames added inline comments.May 5 2021, 5:09 PM
mlir/lib/IR/BuiltinTypes.cpp
516

Maybe we need to talk about this more then.

I think the premise that this is a shape-centric interface only holds for what we have with builtin types.

With custom types, this is not limited to shapes. For example, it would be useful for dialects that want to do type inference with some kind of custom !any_type element type. (It makes me think we probably want to join or meet the element types instead of having a comparison.)

My feeling was the logic will be useful or needed in various places. For example in type inference, or to update the control-flow interfaces be more flexible.

mehdi_amini added inline comments.May 5 2021, 5:26 PM
mlir/lib/IR/BuiltinTypes.cpp
516

There's nothing shape-centric about this interface.

Fair enough: you could imagine this as a type-lattice interface only. But then I'll object that it won't be useful for modeling a shape inference lattice: it would mix on the same lattice properties of the type that looks orthogonal to each other as I see it.

arames updated this revision to Diff 343283.May 5 2021, 9:54 PM

Address review comments.

arames added inline comments.May 5 2021, 9:55 PM
mlir/lib/IR/BuiltinTypes.cpp
503

We may want to do

Type elementTy = joinTypes(ty1.getElementType(), ty2.getElementType());
if (!elementTy)
  return Type();

here and in other places.

silvas added inline comments.May 6 2021, 8:37 AM
mlir/lib/IR/BuiltinTypes.cpp
516

I'm pretty happy with the current design. I view this as a type lattice only, e.g. using it to model that NoneType is a subtype of Optional[T] in Python is what I want to use it for (or refine AnyType to a more specific type).

Indeed, it is not the most useful lattice for shape inference, but that's fine! We can have separate lattice structures for that.

silvas added inline comments.May 6 2021, 8:45 AM
mlir/lib/IR/BuiltinTypes.cpp
503

I don't think we define whether tensor is covariant/contravariant/invariant in the element type (I'm referring to the definitions here, which I think are pretty common in type theory: https://www.python.org/dev/peps/pep-0483/#covariance-and-contravariance). For example, if I have a tensor<optional<T>> is tensor<None> a valid subtype? I don't think we want to guarantee that (it would have implications for bufferization that I would have to think through more fully, for example), so making tensor invariant on the element type would be best for now I think.

jpienaar added inline comments.
mlir/include/mlir/IR/TypeUtilities.h
70

This is the inverse of Shape_JoinOp (and joinShapes was defined as), we went for least generalized there. Lets keep these consistent.

80

Without seeing the use, it is unclear whether populating or returning a SmallVector is most useful.

mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
48

What happens with element type mismatches? Also does this extend to dialects with element types with subtypes?

76

i8?

81

Why is this not tensor<*xf32> ?

mlir/lib/IR/BuiltinTypes.cpp
516

I think this is very good question on encoding. E.g., is no encoding less constrained than tensor with encoding. If so, then join/meet with encoding or without has an answer.

silvas added inline comments.May 6 2021, 3:22 PM
mlir/include/mlir/IR/TypeUtilities.h
70

But DataFlowAnalysis.h has the opposite convention (that's why I liked this one), and seems more centrally positioned.

mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
48

See below about my notes on tensor being covariant/contravariant/invariant on the element type. I think it makes sense to be invariant.

In the case of element type mismatches the result is null because the native tensor type has no intrinsic way to represent "any dtype is ok".

mlir/lib/IR/BuiltinTypes.cpp
516

As I described above, I think that "no encoding" is itself an encoding distinct from any explicitly specified encoding.

silvas added inline comments.May 6 2021, 3:26 PM
mlir/include/mlir/IR/TypeUtilities.h
70

I agree we should keep it consistent. We should either rename the functions in DataFlowAnalysis.h or Shape_JoinOp.

arames marked an inline comment as done.May 6 2021, 3:45 PM
arames added inline comments.
mlir/include/mlir/IR/TypeUtilities.h
70

Although I am accustomed to the naming in this patch, it seems to me changing the internals is easier than changing the op.

80

At this point it is of course only used in this patch.
I would think passing the SmallVectorImpl is more efficient, but I the return-by-value version easier to use.

Do you prefer

LogicalResult joinShapes(ArrayRef<int64_t> shape1, ArrayRef<int64_t> shape2, SmallVectorImpl<int64_t> &join);

?

mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
81

tensor<*xf32> is less specialized than both tensor<4x5xf32> and tensor<6xf32>.
The join (with the current conventions) would yield tensor<*xf32>.

arames added inline comments.May 6 2021, 3:48 PM
mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
48

Couldn't some dialects have a way to represent it though ?
This relates to https://llvm.discourse.group/t/rfc-tensors-with-unknown-element-types/1039.

arames added inline comments.May 6 2021, 3:51 PM
mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
48

(Reading your other comments about this.)

silvas added inline comments.May 6 2021, 4:07 PM
mlir/include/mlir/IR/TypeUtilities.h
80

I think returning by value is fine here. NRVO should do a great job of optimizing this I think: https://www.fluentcpp.com/2016/11/28/return-value-optimizations/

mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
48

I've changed my mind about this since writing that (thanks Chris Lattner for pushing back on it!). I have been working in npcomp using tensor<...x!numpy.any_dtype> and it is a big pain to use, and I don't feel it is worth standardizing. I will be defining my own !torch.tensor type to deal with it (plus there are other considerations like mutability there which I can roll in), and then I can properly define my join/meet functions there.

My experience is that tensor type is mostly useful for ranked tensors of known element type. The unranked version happens to be useful for modeling certain concepts in TensorFlow (which always knows the element type due to how GraphDef work, which is a happy coincidence), but in my experience has very limited utility there, and even less utility in situations where the element type might be unknown, such as numpy/torch/etc. which generally have other considerations, such as mutabilty.

Or let me say it like this. I'm not aware of any situation where an immutable tensor of "unknown element type" that is modeled with a special "!any_element_type" element type feels like the right design. It's more of a "happens to work" design point, which is below the bar for MLIR core.

81

Can you update the code?

81

s/code/comment/

mehdi_amini added inline comments.May 6 2021, 4:13 PM
mlir/include/mlir/IR/TypeUtilities.h
80

The only advantage of the output parameter is the pattern where you can reuse the same container over and over on the caller side to avoid reallocating memory.
(I'm not saying it'll have an impact here)

arames added inline comments.May 6 2021, 4:25 PM
mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
81

???

I was trying to say the current example is fine.
With the current naming conventions,

join(tensor<4x5xf32>, tensor<6xf32>) == tensor<*xf32>
meet(tensor<4x5xf32>, tensor<6xf32>) == Type()
arames added inline comments.May 6 2021, 4:34 PM
mlir/include/mlir/IR/TypeUtilities.h
70

I have not looked much yet at DataFlowAnalysis files. Isn't "join" term used generically to indicate the combination of the states ? The actual implementation of it depends on what the analysis does. It would be "nice" if it matches the "join" used for join/meet type interface, but it does not have to, does it ?
In that case I could only update this patch to match the shape ops.

I don't feel I am in a position to make the decision here.

@silvas @jpienaar

silvas added inline comments.May 6 2021, 4:57 PM
mlir/include/mlir/IR/TypeUtilities.h
70

It doesn't have a super strict dataflow formalism, but it implements a pessimistic* sparse dataflow analysis where getPessimisticValueState is "top" and "join" is least-upper-bound.

*note: it has a hardcoded optimistic treatment of backedges proven unreachable.

I think that both renamings should be easy. For shape.join it's literally just sed'ing shape.join to shape.meet. For DataFlowAnalysis it is renaming that join method.

From a quick survey, my favorite dataflow analysis reference sources both use "top of the lattice is the full set, bottom of the lattice is the empty set" behavior, which corresponds to your current join/meet conventions, so I prefer that, since those are the resources that I tend to point newbies at.

https://www.seas.harvard.edu/courses/cs252/2011sp/slides/Lec02-Dataflow.pdf
https://cs.au.dk/~amoeller/spa/3-lattices-and-fixpoints.pdf

arames updated this revision to Diff 376017.Sep 29 2021, 1:14 PM

Rebase on parent patches.

rriddle added inline comments.Oct 11 2021, 7:24 PM
mlir/include/mlir/Interfaces/JoinMeetTypeInterface.td
55–56
mlir/lib/IR/BuiltinTypes.cpp
528–529

Doesn't this ignore the encoding aspect?

mlir/lib/IR/TypeUtilities.cpp
161–163
178–180
mlir/lib/Interfaces/JoinMeetTypeInterface.cpp
20–21

Can we just assert this? (Although, I'd honestly just not check this at all)

27–29

Can we avoid computing the second join in opt builds if we can (i.e. if the first join succeeded, only compute the second inside of the assert)? (You could also realistically compare the TypeIDs of the two types and avoid the second join altogether).

40–41

Same here.

43–53

Same here.

arames marked 22 inline comments as done.Oct 13 2021, 10:56 AM
arames added inline comments.
mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
81

Why is this (tensor<1xi32> ^ i32) not tensor<*xf32> ?

I don't think we want scalar and tensor types to join/meet on the type lattice.
Why would i32 be more specialized than tensor<i32>, and not the opposite ?

We would not expect an op accepting tensors only could not take an i32 where a tensor<*xi32> is expected.

mlir/lib/IR/BuiltinTypes.cpp
516

Yes. Now fixed.

528–529

Doesn't this ignore the encoding aspect?

Absolutely. I had not addressed the earlier comments yet.
This is not fixed, with the associated test

"meet"(%tensor_1xi32_encoding1, %tensor_unranked_i32) : (tensor<1xi32, #encoding1>, tensor<*xi32>) -> i1
// expected-error@-1 {{types do not meet}}
mlir/lib/Interfaces/JoinMeetTypeInterface.cpp
20–21

Do you mean assert(ty1 && ty2) ?

I think we explicitly want to allow null input types, because it allows trivially chaining calls to joinTypes and meetTypes without having to check for null. It

27–29

I like the idea.
Changing to do that. Others can comment if they disagree.

arames updated this revision to Diff 379468.Oct 13 2021, 10:56 AM
arames marked 4 inline comments as done.

Address review comments:

  • Fix join/meet for tensors with encodings
  • misc style fixes
arames added inline comments.Oct 13 2021, 10:58 AM
mlir/include/mlir/Interfaces/JoinMeetTypeInterface.h
81

This should have been:
We would not expect an op accepting tensors only could take an i32 where a tensor<*xi32> is expected.

arames updated this revision to Diff 382781.Oct 27 2021, 1:55 PM

Rebase on top of tree.

arames updated this revision to Diff 382782.Oct 27 2021, 1:57 PM

Fix clang-format warning.

arames updated this revision to Diff 383879.Nov 1 2021, 1:53 PM

Rebase on top of tree.