Page MenuHomePhabricator

[mlir][bufferization] Add OneShotBufferize transform op
ClosedPublic

Authored by springerm on May 6 2022, 8:07 AM.

Details

Summary

This commit allows for One-Shot Bufferize to be used through the transform dialect. No op handle is currently returned for the bufferized IR.

Diff Detail

Event Timeline

springerm created this revision.May 6 2022, 8:07 AM
Herald added a project: Restricted Project. · View Herald TranscriptMay 6 2022, 8:07 AM
springerm requested review of this revision.May 6 2022, 8:07 AM
Herald added a project: Restricted Project. · View Herald TranscriptMay 6 2022, 8:07 AM

Nice, thank you!

I have one question about the targeting model. If the op requires some payload IR module to be passed as target (which also removes it from the mapping rendering it useless for further transformations), I feel like it should also return a handle to the bufferized module, otherwise we will have hard time composing with any transforms after the bufferization.

mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
24

Nit: it's not very clear what actually gets bufferized given the target. Isn't bufferization a module-level transform meaning target should always be a module op?

31
39–40

I'm afraid you need to manually implement the MemoryEffectsOpInterface to also indicate a read/write effect on PayloadIRResource.

springerm added inline comments.May 13 2022, 3:09 AM
mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
24

It is often applied to an entire module, but it can be applied to arbitrary operations. (Maybe the old bufferization did not allow this.)

39–40

Now that I'm looking at this again, this is actually not as precise as it could be.

If target is a ModuleOp, it survives bufferization. It is modified in-place. (Or to be precise, its regions are modified, but the op stays the same. This guaranteed and should be documented.) This is not necessarily the case for other ops. They may be updated in-place, but most likely they are deleted and replaced with something else.

But I think I cannot express this in TileOp::getEffects. I would have to check the type of target. If it is a ModuleOp, there should be no TransformMappingFree effect.

Adding a return value is difficult because there may not be one for certain ops. Maybe they disappear entirely for some reason. (A bufferization pattern could choose to do this.) If even they do not disappear, I currently do not get a handle to the new bufferized op from the bufferization.

Alternatively, we could also just limit this transform op to ModuleOps. Then we also don't need a return value because the handle can be reused. (Or we just return a new handle to the same ModuleOp.)

What do you think?

ftynse added inline comments.May 19 2022, 8:42 AM
mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
39–40

We can also have two different ops or an attribute on this op to switch between the two cases:
(1) bufferizing the entire module, the handle is not freed
(2) bufferizing a piece of a module, the handle is freed (but the handle to the module is not, which lets us pdl_match again after bufferization).

springerm updated this revision to Diff 435490.Thu, Jun 9, 3:57 AM
springerm marked 2 inline comments as done.

address comments

springerm added inline comments.Thu, Jun 9, 3:59 AM
mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
39–40

added an attribute

ftynse accepted this revision.Thu, Jun 9, 5:50 AM
This revision is now accepted and ready to land.Thu, Jun 9, 5:50 AM
This revision was automatically updated to reflect the committed changes.