Page MenuHomePhabricator

[mlir][RFC] Refactor layout representation in MemRefType
ClosedPublic

Authored by vinograd47 on Oct 11 2021, 8:32 AM.

Details

Summary

The change is based on the proposal from the following discussion:
https://llvm.discourse.group/t/rfc-memreftype-affine-maps-list-vs-single-item/3968

  • Introduce MemRefLayoutAttr interface to get AffineMap from an Attribute (AffineMapAttr implements this interface).
  • Store layout as a single generic MemRefLayoutAttr.

This change removes the affine map composition feature and related API.
Actually, while the MemRefType itself supported it, almost none of the upstream
can work with more than 1 affine map in MemRefType.

The introduced MemRefLayoutAttr allows to re-implement this feature
in a more stable way - via separate attribute class.

Also the interface allows to use different layout representations rather than affine maps.
For example, the described "stride + offset" form, which is currently supported in ASM parser only,
can now be expressed as separate attribute.

Diff Detail

Event Timeline

vinograd47 created this revision.Oct 11 2021, 8:32 AM
vinograd47 requested review of this revision.Oct 11 2021, 8:32 AM
Herald added a project: Restricted Project. · View Herald Transcript
bondhugula requested changes to this revision.Oct 11 2021, 10:00 AM
bondhugula added inline comments.
mlir/include/mlir/IR/BuiltinTypes.td
575–576

Missing doc comments here.

mlir/include/mlir/IR/MemRefLayoutAttrInterfaces.td
18–24

Is this attribute interface missing its documentation? It's crucial to document this.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
86–89

All this code is going to be unsafe and incorrect since the getAffineMap() can now return null given this change.

mlir/lib/IR/BuiltinTypes.cpp
659–662

This would mean that every piece of code getting an affine map from the memref type would have to check for non-null, and if not handling it, propagate the failure/error all the way. What kind of arbitrary layout you would like to support that can't be currently represented by an affine map and that is not better served by a custom dialect-specific memref type? In other words, shouldn't the builtin MemRefType always have a layout attribute that implements a MemRefLayoutAttrInterface?

This revision now requires changes to proceed.Oct 11 2021, 10:00 AM

I am very supportive of this direction. It contributes towards MLIR generality and uses the "little builtin" design principle. It also reduces the conceptual hack of AffineMap being a separate object from AffineMapAttr (AFAIK, MemRef was the only place that used pure AffineMap rather than have it wrapped in an attribute).

A couple of high-level comments, I'll do another pass with in-depth review later.

mlir/include/mlir/IR/BuiltinTypes.td
575

Since we are changing the API already, I would consider not providing these fallbacks. Changing getAffineMaps().front() to getLayout().getAffineMap() and getAffineMaps().front().isIdentity() to getLayout().isIdentity() isn't a huge change, modulo the null issue @bondhugula raises.

mlir/include/mlir/IR/MemRefLayoutAttrInterfaces.h
1

There's already https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/BuiltinTypeInterfaces.td, let's just put this interface there as well.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
86–89

Out of curiosity, are you aware of cases that actually use symbol operands in load/store operations beyond unit tests?

mlir/lib/IR/BuiltinTypes.cpp
655

I think this should return MemRefLayoutAttInterface instead of generic Attribute.

659–662

+1 to MemRefType always having an attribute that implements MemRefLayoutAttrInterface.

One scenario that is not expressible using affine maps to my knowledge is sparse storage buffers. We support sparse layouts on tensors but have no buffer equivalent and use opaque i8 buffers instead. Another is expressing relations between dimensions or other performance-critical guarantees (divisibility, contiguity) of the layout that we may want to add that are not exactly the layout information, but may be added to the attribute. Yes, it is possible to reimplement a MemRef type each time a custom layout is needed. Beyond the usual concern with code duplication (this concern has been raised for simple operations such as dim, the memref type is significantly more complex), it is all the code built around memref as the only buffer type that will not work on that type. Longer term, we may want a BufferTypeInterface that generalizes the common aspects of types required by the code working on buffers. However, introducing it looks like a more intrusive change than what is proposed here.

I am very supportive of this direction. It contributes towards MLIR generality and uses the "little builtin" design principle. It also reduces the conceptual hack of AffineMap being a separate object from AffineMapAttr (AFAIK, MemRef was the only place that used pure AffineMap rather than have it wrapped in an attribute).

A couple of high-level comments, I'll do another pass with in-depth review later.

+1, thanks for doing this @vinograd47 !

vinograd47 edited the summary of this revision. (Show Details)

Updated according to review remarks

vinograd47 marked 6 inline comments as done.Oct 15 2021, 8:47 AM

@ftynse @bondhugula I've addressed your comments.

As for sparse storage buffers. Well, my personal opinion, they should be represented as separate buffers (data + indexes). Just, IMHO, merging such structure in single MemRef, will overcomplicate it. But again this is just my personal thoughts.

bondhugula added inline comments.Oct 15 2021, 9:49 AM
mlir/include/mlir/IR/BuiltinTypes.td
524–526

With this change to always have MemRefLayoutAttr, this revision and refactoring to the layout interface is looking good to me. Thanks!

533–557

We should avoid this much C++ code inside .td files. Can we move them to a .cpp file?

mlir/lib/IR/BuiltinTypes.cpp
659–662

For future extensions for more generic layouts that aren't captured by affine maps, I would exactly go in the direction you are suggesting in the last two lines (BufferTypeInterface or even MemRefKindInterface). I feel it's in the right direction in all ways and we'd want to avoid two highly invasive changes: trying to generalize the existing MemRefType by adding a vastly different layout variety would lead to all kinds of handling and error propagation logic for the vast majority of buffer use cases. It'd be quite common I think for downstream users to need both custom tensor types and custom memref types while wanting to "inherit" as much as possible for the builtin tensor and memref types.

A list of mostly minor comments.

mlir/include/mlir/IR/BuiltinAttributeInterfaces.td
444 ↗(On Diff #379991)

type -> type's
semantic -> semantics
"..., and offsets."

-> "Such a layout attribute should be representable as a ..."?

456 ↗(On Diff #379991)

the identity layout.

mlir/lib/Analysis/LoopAnalysis.cpp
224–226

Drop trivial braces.

mlir/lib/Analysis/Utils.cpp
618–621

Likewise.

mlir/lib/Dialect/Vector/VectorOps.cpp
3666–3668

I think these error messages need to be updated? But perhaps this was an issue with the pre-existing code itself.

mlir/lib/IR/BuiltinTypes.cpp
695

spici... -> speci

Missing Layout ... -> Missing layout ...

mlir/lib/Transforms/NormalizeMemRefs.cpp
272

This is using the result of a dyn_cast unchecked - a pre-existing issue. Should have been a cast. I can fix this in another revision.

mlir/test/IR/parser.mlir
80

Is this comment outdated now?

88

Likewise.

Addressed review comments

vinograd47 marked 12 inline comments as done.Oct 18 2021, 5:48 AM

@bondhugula I've addressed your comments, please take a look one more time.

mlir/lib/Analysis/Utils.cpp
618–621

It is not a trivial braces case, there are 2 statements in the body (LLVM_DEBUG + return).

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
86–89

Not sure who you addressed this question to. Personally, I don't know the answer, since we don't use load/store operations in our project.

ftynse accepted this revision.Oct 18 2021, 6:08 AM

Some more nits from me, LGTM otherwise.

mlir/include/mlir/IR/BuiltinAttributeInterfaces.h
238 ↗(On Diff #380354)

the the

mlir/include/mlir/IR/BuiltinAttributeInterfaces.td
439 ↗(On Diff #380354)

Nit: I'd call the class MemRefLayoutAttrInterface with the explicit "Interface" word in it. The ElementsAttr above doesn't have it for mostly legacy reasons (it used to be a proper attribute before becoming an interface).

mlir/lib/Bindings/Python/IRTypes.cpp
417–419

Prefer mlirAttributeGetNull to default initialization.

457

Nit: "the layout of the MemRef type as affine map", otherwise it reads like the affine map is a different property.

mlir/lib/CAPI/IR/BuiltinTypes.cpp
233

Let's do something like MlirAttributeIsNull(layout) ? MemRefLayoutAttr() : unwrap(layout).cast<MemRefLayoutAttr>() here and below. Otherwise, we will be accepting wrong layout types and silently ignoring them.

mlir/lib/IR/BuiltinTypes.cpp
659–662

Then I would suggest @vinograd47 to add this conclusion somewhere in the documentation along with this patch.

695

Can we have a test for this message?

Tiny nit: we start error messages with a small letter because they are printed after some other text.

vinograd47 marked 2 inline comments as done.

Addressed review comments

vinograd47 marked 7 inline comments as done.Oct 18 2021, 6:55 AM

@ftynse I've addressed your comments.

mlir/lib/IR/BuiltinTypes.cpp
695

Can we have a test for this message?

Actually NULL layout attributes are replaced with identity affine map in the get methods implementation. I've left this NULL check in the verify method just for extra safety. Not sure if it is possible to test it.

ftynse added inline comments.Oct 18 2021, 7:03 AM
mlir/lib/IR/BuiltinTypes.cpp
695

Then let's replace it with an assert.

rriddle added inline comments.Oct 18 2021, 7:14 AM
mlir/lib/Analysis/LoopAnalysis.cpp
222–224

Is this TODO resolved now?

vinograd47 marked an inline comment as done.

@ftynse @rriddle Addressed your comments.

bondhugula accepted this revision.Oct 19 2021, 2:07 AM

LGTM - thanks very much.

This revision is now accepted and ready to land.Oct 19 2021, 2:07 AM