This is an archive of the discontinued LLVM Phabricator instance.

[mlir] Add BufferResultsToOutParams pass.
ClosedPublic

Authored by silvas on Oct 23 2020, 12:53 PM.

Details

Summary

This pass allows removing getResultConversionKind from
BufferizeTypeConverter. This pass replaces the AppendToArgumentsList
functionality. As far as I could tell, the only use of this functionlity
is to perform the transformation that is implemented in this pass.

Future patches will remove the getResultConversionKind machinery from
BufferizeTypeConverter, but sending this patch for individual review for
clarity.

Diff Detail

Event Timeline

silvas created this revision.Oct 23 2020, 12:53 PM
silvas requested review of this revision.Oct 23 2020, 12:53 PM
dfki-mako added inline comments.Oct 27 2020, 4:41 PM
mlir/include/mlir/Transforms/Passes.td
220

Just for my understanding: This pass basically simplifies the use of populateWithBufferizeOpConversionPatterns via a custom Rewriter instance for an input dialect, right? Do you plan to remove the conversion patterns in favor of this pass in the future?

My question is here: Why should this pass be executed after the entire bufferization process? Isn't it useful to apply this pass before applying the BufferDeallocation pass, which inserts all the necessary deallocations? Or do you refer to the first part of the "overall" bufferization process, which ports the input program into the "buffer world" (without taking deallocations into account)?

Regarding the Linalg dialect dependency: We are going to discuss possible solutions to remove the Std- and Linalg dialect dependencies from the BufferDeallocation pass in the public Discourse-Forums these days. I suspect that the result of this discussion can also be applied to this pass (in the future).

silvas added inline comments.Oct 27 2020, 5:00 PM
mlir/include/mlir/Transforms/Passes.td
220

This pass eliminates the need for the AppendToArgumentList conversion kind, which allows removing that complexity from BufferizeTypeConverter. (and together with removing the tuple stuff, BufferizeTypeConverter can just be a regular type converter; everything gets quite a bit simpler).

Regarding BufferDeallocation, I think we agree but just got confused by terminology. "bufferization" (as defined in Bufferize.h and being used here) is about applying conversions. So this would be applied right after applying the conversions, but before any buffer optimizations.

I guess this approach introduces a limitation with respect to applicability to arbitrary dialects. We are now able to match ReturnOp and CallOps only while the previous version supported a generic customization to arbitrary dialects, right?
I would recommend that we should try to match CallOpInterface and ReturnLike implementations to make this code more generic in favor of standard CallOp and ReturnOp-based matchers from the Std dialect.

mlir/include/mlir/Transforms/Passes.td
220

AppendToArgumentList conversion kind, which allows removing that complexity from BufferizeTypeConverter

Sounds reasonable, since it simplifies the type converter to become "just a type converter" (as you said).

Regarding BufferDeallocation, I think we agree but just got confused by terminology. "bufferization"

Sounds pretty good to me, too. In this context, this pass should be the last one to be executed (as you mentioned) to "fix" all return ops and call sites (at least for statically shaped memrefs).

mlir/lib/Transforms/BufferResultsToOutParams.cpp
37

functionType (above) vs func.getType() (2 uses)

43

I am note sure about the braces here since it is a single statement. However, they improve readability in this case.

52

Suggestion: maybe we could change the function to receive an additional vector by reference and push the results into this vector instead of returning an empty vector in this case.

58

nit: maybe a tiny comment would be nice in this case
something like: // Updates all ReturnOps in the scope of the given FuncOp by either keeping them as return values or copying the associated buffer contents into the given out params.

70

Suggestion: maybe change t to outParamEntry. Regarding the braces: same comment as above.

79

nit: maybe another tiny comment would be nice in this case, as well
something like: // Updates all CallOps in the scope of the given ModuleOp by allocating temporary buffers for newly introduced out params

93

Maybe I am missing the point but can't we just check this in the first loop to avoid this additional cast?

111

nit: maybe change t to something more meaningful.

126

Suggestion: If we change the signature of the updateFuncOp function, we can avoid this variable. (see above)

mlir/test/Transforms/buffer-results-to-out-params.mlir
97

As you are updating the tests anyway, it would be really nice to break after 80 chars.

dfki-mako requested changes to this revision.Oct 28 2020, 5:42 AM
This revision now requires changes to proceed.Oct 28 2020, 5:42 AM
silvas updated this revision to Diff 301444.Oct 28 2020, 2:52 PM
silvas marked 7 inline comments as done.

Address @dfki-mako comments.

Thanks Marcel for the very thorough review!

I guess this approach introduces a limitation with respect to applicability to arbitrary dialects. We are now able to match ReturnOp and CallOps only while the previous version supported a generic customization to arbitrary dialects, right?
I would recommend that we should try to match CallOpInterface and ReturnLike implementations to make this code more generic in favor of standard CallOp and ReturnOp-based matchers from the Std dialect.

Good point! I really would like to write this pass as being generic as well. I tried to write this pass against CallOpInterface and ReturnLike and FunctionLike, but they currently don't have enough functionality for this transformation, in particular the parts where we need to mutate/recreate the ops.

For example, FunctionLike doesn't guarantee that the op has a FunctionType, so all the code here where we modify the function type can't be written generically (see this thread for a similar issue:
https://llvm.discourse.group/t/moving-erasearguments-from-funcop-to-functionlike/2016/14). In particular, even after the patches in that thread, we cannot append new result types. (also, I don't think that traits like FunctionLike / ReturnLike.

Also, note that FunctionLike is a trait, not an interface, so it cannot be accessed dynamically. One must already have an object of the concrete type to use it.

mlir/lib/Transforms/BufferResultsToOutParams.cpp
93

I conceptually associate this error with the creation of the AllocOp -- we don't care about this limitation elsewhere in the entire pass, and one can imagine ways to avoid this, which would involve changing this code, rather than the code above (which merely determines which results we need to convert to an out param -- not how we do it).

For example, we could require annotating functions with an attribute that somehow describes how to compute their out param sizes e.g. func @callee(...) -> (...) attributes {shape_transfer_function = @compute_callee_out_param_shapes}, or do so via some analysis, etc. In that circumstance, the static shape case is but one of multiple ways that we would know how to allocate the results.

Also this pass only runs once in the pipeline, so performance isn't that critical. For example, we are walking every op in the module just to find call ops, and we effectively do the same for return. If we really cared about performance, we would only do a single walk of the whole module, but that would make the code a lot more confusing.

111

Same comment as above, I think that succinctness here makes the code clearer, as the minimum meaningful name is just too long and distracting.

It's too bad that C++ doesn't let us do something like this:

for (auto [replaceMe, newValue] : llvm::zip(replaceWithNewCallResults, newCall.getResults()))
  replaceMe.replaceAllUsesWith(newValue)

I think that C++17 has a feature for this, but I don't know if we allow using it yet: https://en.cppreference.com/w/cpp/language/structured_binding

mlir/test/Transforms/buffer-results-to-out-params.mlir
97

I typically don't break in test cases so that it reflects the typical printed IR structure. The theory is that our brains are mostly tuned to parse non-line-broken IR since that is what we usually read (unlike in C++), so line-breaking it tends to be more confusing since there aren't clear style preferences and mental heuristics for parsing it.

Also, to answer your specific question, the previous version did not support a generic customization to arbitrary dialects. The FuncOp / CallOp patterns in the previous code are hard-coded (which as I described, cannot be removed with the current set of traits/interfaces). Only BufferizeReturnOpConverter converter is templatized, but without being able to convert generic function-like ops, that doesn't seem to be very useful (since converting the function-like ops is the hardest part).

For example, FunctionLike doesn't guarantee that the op has a FunctionType, so all the code here where we modify the function type can't be written generically (see this thread for a similar issue:

Unfortunately, we cannot use this trait for the reasons you mentioned.

Also, to answer your specific question, the previous version did not support a generic customization to arbitrary dialects.

Yeah, you are absolutely right. Consequently, we will not lose any genericity by merging your CL, which is really great +1
Future extensions in the MLIR world might enable us to make this pass more generic with respect to FuncOps.

dfki-mako added inline comments.Oct 29 2020, 3:55 PM
mlir/lib/Transforms/BufferResultsToOutParams.cpp
93

Sounds reasonable. I guess we can safely keep it this way +1.

111

Yes, unfortunately that is also true. As far as I know, we shouldn't use this extension at the moment.

mlir/test/Transforms/buffer-results-to-out-params.mlir
97

+1

Awesome work +1. I guess this should be ready to go after replacing some autos with their underlying types.

mlir/lib/Transforms/BufferResultsToOutParams.cpp
66

nit: auto -> OpOperand

88

nit: auto -> OpResult

96

nit: auto -> Value

silvas updated this revision to Diff 301792.Oct 29 2020, 5:21 PM
silvas marked 2 inline comments as done.

address comments

dfki-mako accepted this revision.Oct 29 2020, 5:38 PM
This revision is now accepted and ready to land.Oct 29 2020, 5:38 PM
This revision was landed with ongoing or failed builds.Oct 30 2020, 2:08 PM
This revision was automatically updated to reflect the committed changes.