Page MenuHomePhabricator

[mlir][bufferize] Add argument materialization for bufferization
ClosedPublic

Authored by herhut on Nov 24 2020, 6:57 AM.

Details

Summary

This enables partial bufferization that includes function signatures.

Diff Detail

Event Timeline

herhut created this revision.Nov 24 2020, 6:57 AM
herhut requested review of this revision.Nov 24 2020, 6:57 AM
silvas added inline comments.Nov 24 2020, 9:08 AM
mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp
34

FuncOpTypeConversionPattern will change the type of block arguments, which requires that all predecessor successor args are converted with a pattern like you have for ReturnOp.

See the rationale here: https://github.com/llvm/llvm-project/blob/25777080549bb62b6e46a1809f93257969f5dd53/mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td#L37

You will need to add a pattern that generically converts all BranchOpInterface ops and set up legality so that if we have a terminator that doesn't implement BranchOpInterface then conversion fails.

Once we do that, I think we can change func-bufferize to be a partial bufferization only, and have a "finalizing-bufferize" that just folds tensor_load/tensor_to_memref.

silvas added inline comments.Nov 24 2020, 9:28 AM
mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp
34

(This is actually how I wanted to originally design it (and it is a better design IMO), but I was too lazy / couldnt't figure out how to do the BranchOpInterface rewrite / legality; in retrospect I shouldn't have taken that shortcut)

silvas added inline comments.Nov 24 2020, 10:11 AM
mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td
41

I was never happy with this "just documented" restriction. Using the design I described below, we can improve this situation, because then the finalizing bufferize is just folding tensor_to_memref/tensor_load which we can do while ensuring that no invalid IR is created!

herhut updated this revision to Diff 307653.Nov 25 2020, 10:24 AM

Rework to support terminators that implement interfaces.

herhut marked 2 inline comments as done.Nov 25 2020, 10:31 AM

PTAL.

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp
34

How about this version? If this works for you, I'll fix the documentation and add another test. Just wanted to get your opinion pre-thanksgiving so that I can go ahead with this (and unblock some other work).

We might still want to keep the tests and pass for the finalizing version (to make sure that use case works, as well) as many users might bufferize that way.

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
88

A cleanup would be to implement the interfaces or have a new one for "end of list" terminators. That's more a cleanup to generalize this, though.

I'm okay to submit this if you make func-bufferize only be a partial conversion (never a finalizing conversion), and add a separate finalizing-bufferize pass which folds away tensor_load/tensor_to_memref and emits an error if any remain.

Note that in that new world, mlir-opt -func-bufferize -finalizing-bufferize will reproduce the behavior of the current func-bufferize, so it is trivial for users to update; also, I wrote all the current users of func-bufferize so there is no need to worry about any other users ;)

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
91

This should just return failure() (or rewriter.notifyMatchFailure(op, "terminator does not implement required interfaces"))

The reason is that patterns should not take an opinion on emitting errors in cases like these. The errors should naturally come from how the conversion has the legality set up. For example, if a dialect conversion pass also adds a special pattern for a specific terminator, we don't want to be emitting errors from this pattern.

mlir/lib/Transforms/Bufferize.cpp
20

would prefer to call it materializeTensorLoad and for it to be a local variable in BufferizeTypeConverter::BufferizeTypeConverter just so it is closer to where it is used.

mlir/test/Dialect/Standard/func-bufferize-partial.mlir
2

nit: is --debug-only needed for the test?

herhut updated this revision to Diff 307790.Nov 26 2020, 1:39 AM

Split out finalizing pass.

rriddle added inline comments.Nov 26 2020, 1:49 AM
mlir/include/mlir/Transforms/Passes.td
622

These should be sortsd.

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp
54

IsKnownNonTerminator

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
47

Comments should be ///

59

Please use the proper terminator checks, isKnownNonTerminator

herhut marked 4 inline comments as done.

Frederik, can you take a look? Sean was happy for this to move forward with his comments addressed, which I did.

mlir/lib/Transforms/Bufferize.cpp
20

It is beyond my C++ voodoo skills to make that happen. If I use a local variable then the compiler complains

ERROR: third_party/llvm/llvm-project/llvm/include/llvm/ADT/STLExtras.h:118:59 type '(lambda at third_party/llvm/llvm-project/mlir/lib/Transforms/Bufferize.cpp:32:32) &' cannot be used prior to '::' because it has no members
struct function_traits : public function_traits<decltype(&T::operator())> {};

Using a static fixes this :)

frgossen accepted this revision.Nov 26 2020, 2:08 AM

Sorry, I saw this only now.
Thank you!

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp
41–42

I've seen this a couple of times now. Maybe it is worth moving this to the type converter.

56

nit: trivial braces

mlir/test/Dialect/Standard/func-bufferize-partial.mlir
17

nit: // ----- missing?

mlir/test/Dialect/Standard/func-bufferize.mlir
1

nit: Some are -, some are --

This revision is now accepted and ready to land.Nov 26 2020, 2:08 AM
rriddle added inline comments.Nov 26 2020, 2:11 AM
mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp
41–42

It already exists, just pass the op itself.

herhut updated this revision to Diff 307807.Nov 26 2020, 3:00 AM
herhut marked 10 inline comments as done.

Comments

Thanks for the reviews!

mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp
54

This does not work for unknown operations. I have changed it now to first check isKnownNonTerminator to handle known operations efficiently and then also check for whether it is the last op in the block to handle unknown operations.

mlir/test/Dialect/Standard/func-bufferize-partial.mlir
17

It is not required, only for the last one as it will abort the conversion.

rriddle added inline comments.Nov 26 2020, 3:04 AM
mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp
54

!isKnownTerminator?

herhut added inline comments.Nov 26 2020, 3:16 AM
mlir/lib/Dialect/StandardOps/Transforms/FuncBufferize.cpp
54

If I use !isKnownTerminator then this would declare unknown operations in terminator position as legal. But I want to fail the lowering if there is a block with an unknown terminator.

herhut updated this revision to Diff 307829.Nov 26 2020, 4:26 AM

Do not rewrite ReturnLike ops in general.

I realized that we cannot rewrite all ReturnLike operations, as they require rewriting of their parent which cannot be done generically. PTAL the changes.

frgossen accepted this revision.Nov 26 2020, 4:35 AM

Thanks!