This is an archive of the discontinued LLVM Phabricator instance.

[mlir][IR] Trigger notifyOperationRemoved callback for nested ops
ClosedPublic

Authored by springerm on Feb 16 2023, 8:03 AM.

Details

Summary

When cloning an op, the notifyOperationInserted callback is triggered for all nested ops. Similarly, the notifyOperationRemoved callback should be triggered for all nested ops when removing an op.

Listeners may inspect the IR during a notifyOperationRemoved callback. Therefore, when multiple ops are removed in a single RewriterBase::eraseOp call, the notifications must be triggered in an order in which the ops could have been removed one-by-one:

  • Op removals must be interleaved with notifyOperationRemoved callbacks. A callback is triggered right before the respective op is removed.
  • Ops are removed post-order and in reverse order. Other traversal orders could delete an op that still has uses. (This is not avoidable in graph regions and with cyclic block graphs.)

Diff Detail

Event Timeline

springerm created this revision.Feb 16 2023, 8:03 AM
Herald added a project: Restricted Project. · View Herald Transcript
springerm requested review of this revision.Feb 16 2023, 8:03 AM

Is this meant to fix something?

I have a use case for it. I could work around it. But it's more of a consistency issue. (We notify of nested ops when they are created, so we should also notify when they are removed.) Not many people seem to be using the listeners, so I think nobody noticed this until now.

springerm planned changes to this revision.Feb 16 2023, 10:54 AM

The GreedyPatternRewriteDriver can be simplified in the same go.

Also simplify GreedyPatternRewriteDriver as part of this change.

rriddle accepted this revision.Feb 16 2023, 11:53 PM

How does the interact with dialect conversion?

This revision is now accepted and ready to land.Feb 16 2023, 11:53 PM

How does the interact with dialect conversion?

Dialect conversion bypasses this notification entirely. It overrides eraseOp and replaceOp. (There may be some opportunities to reuse a larger part of the RewriterBase infrastructure in various places of the dialect conversion.)

change traversal order: bottom to top

@rriddle There's an issue that I overlooked. The order in which the notifications should be triggered. It should be bottom-to-top. Because ops can only be removed bottom-to-top. And we should trigger the callback in the same order as if the ops were removed one-by-one. Does this sounds reasonable to you? I just noticed this because I had failures in the GreedyPatternRewriteDriver.

I think we don't have a way to walk ops bottom-to-top at the moment. Thus the custom traversal. I'm going to prepare a change that adds another parameter to Operation::walk and rebase this revision on top of it.

springerm updated this revision to Diff 498339.Feb 17 2023, 5:24 AM
springerm edited the summary of this revision. (Show Details)

rebase

Mogball accepted this revision.Feb 19 2023, 12:32 AM
springerm updated this revision to Diff 499841.Feb 23 2023, 6:46 AM

Add test case

springerm updated this revision to Diff 499875.Feb 23 2023, 8:33 AM

interleave notifications and op removal

springerm edited the summary of this revision. (Show Details)Feb 23 2023, 8:40 AM

@Mogball @rriddle @mehdi_amini I made some significant changes to this revision since it was accepted. Could you take another look?

mehdi_amini added inline comments.Feb 23 2023, 3:07 PM
mlir/lib/IR/PatternMatch.cpp
292

Technically this return true is valid only for registered dialect. But seems like this is a debug function so not that important I guess

320

Blocks are not necessarily sorted topologically, how do you handle that?

springerm added inline comments.Feb 24 2023, 3:33 AM
mlir/lib/IR/PatternMatch.cpp
320

I didn't know that such IR was valid.

I think ideally we would walk the IR by following the "reverse successor" relationship of blocks. Not sure if it is really necessary here, there may be a simpler solution. Definitely something needs to be changed because I found an example that crahes.

Related to this: I've seen a use case in the DialectConversion where IR must be walked according to the "forward successor" relationship.

We could utilize the new Iterator template of Operation::walk: By providing two new iterators: ForwardDominanceIterator and ReverseDominanceIterator. They work like ForwardIterator and ReverseIterator but walk blocks according to successors/reverse successors.

These iterators would be more heavyweight than simple forward/reverse iterators (i.e., they need to keep track of visited blocks etc. in case of cycles). Also, the current detail::walk implementations would have to be moved from Visitors.cpp to Visitors.h, which is already tricky without including Operation.h in Visitors.h (cyclic include).

fix unstructured control flow

springerm added a comment.EditedMar 15 2023, 10:15 AM

This should now handle arbitrary successor/predecessor structures and/or graph regions correctly.

springerm planned changes to this revision.Mar 16 2023, 12:46 AM

This does not work correctly if an op in a dead region has nested ops. I was trying to do without, but we really need a ReverseDominanceIterator that also enumerates dead blocks.

springerm updated this revision to Diff 507286.Mar 22 2023, 2:53 AM

Perform a custom iteration over regions/blocks/ops. This is much simpler than trying to make ReverseDominanceIterator enumerate dead (unreachable) blocks.

This revision is now accepted and ready to land.Mar 22 2023, 2:53 AM
springerm marked an inline comment as done.Mar 22 2023, 2:55 AM
springerm added inline comments.
mlir/lib/IR/PatternMatch.cpp
320

This is now fixed by iterating blocks post-order.

LGTM, but @rriddle should have another look

springerm updated this revision to Diff 508290.Mar 25 2023, 2:48 AM

rebase and trigger new CI run

springerm edited the summary of this revision. (Show Details)Mar 27 2023, 12:12 AM
springerm updated this revision to Diff 508558.Mar 27 2023, 3:39 AM
springerm edited the summary of this revision. (Show Details)

also drop uses of bbArg