Page MenuHomePhabricator

[mlir] Extend Operation visitor with pre-order traversal
ClosedPublic

Authored by dcaballe on Feb 22 2021, 12:27 PM.

Details

Summary

This patch extends the Region, Block and Operation visitors to also support pre-order walks.
We introduce a new template argument that dictates the walk order (only pre-order and
post-order are supported for now). The default order for Regions, Blocks and Operations is
post-order. Mixed orders (e.g., Region/Block pre-order + Operation post-order) could easily
be implemented, as shown in NumberOfExecutions.cpp.

Diff Detail

Event Timeline

dcaballe created this revision.Feb 22 2021, 12:27 PM
dcaballe requested review of this revision.Feb 22 2021, 12:27 PM
dcaballe added inline comments.Feb 22 2021, 12:30 PM
mlir/lib/IR/Visitors.cpp
43

Walk implementation for Region and Block callbacks above seems to be already happening in pre-order. Is this expected? I found this a bit confusing since the traversal for Operation callbacks is in post-order by default.

  1. We could change the default traversal for Region and Block callbacks to post-order, if that made any sense, and let them take the IsPreOrder argument.
  1. Otherwise, we could keep the current behavior and separate the public Region/Block walk implementation in Visitor.h from the Operation implementation so that only the latter takes the IsPreOrder template argument.

WDYT?

Looks mostly good. My main comment is that I would not use a boolean for the template argument, it is very ambiguous what true/false means. Can you use something like an enum instead?

rriddle added inline comments.Feb 22 2021, 2:26 PM
mlir/lib/IR/Visitors.cpp
43

I would ensure that everything has the same default behavior. So I think option 1 (If I'm reading that right).

Please also see: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/utils/visitor_util.h

Thanks for the pointer! It would be great to have something as generic as that visitor in the mlir repo!

Can you use something like an enum instead?

Sure!

mlir/lib/IR/Visitors.cpp
43

Yes, that's #1. I'll give it a try. Hopefully, changing the current default behavior doesn't require too many changes.

frgossen accepted this revision.Feb 23 2021, 12:48 AM

Very useful. Thanks!

mlir/lib/IR/Visitors.cpp
43

I found this strange, too. I believe the liveness analysis, or at least its tests, depend on this behaviour.

#1 sounds like the way to go.

This revision is now accepted and ready to land.Feb 23 2021, 12:48 AM
dcaballe updated this revision to Diff 325921.Feb 23 2021, 3:47 PM

Addressing feedback:

  • Use enum to encode walk order.
  • Set region, block and operation walks to post-order by default.
  • Adjust walks in Liveness and NumberOfExecutions to align with the new default traversal.
dcaballe edited the summary of this revision. (Show Details)Feb 23 2021, 3:48 PM
rriddle accepted this revision.Feb 23 2021, 4:59 PM

LGTM after resolving last comment.

mlir/lib/IR/Visitors.cpp
16

Can you just pass the order as an extra parameter? These methods represent the internal implementation of the walk methods, so I'm fine with adding more parameters here. That would also remove all of the templates within this file.

bondhugula requested changes to this revision.Feb 23 2021, 5:53 PM
bondhugula added a subscriber: bondhugula.
bondhugula added inline comments.
mlir/lib/IR/Visitors.cpp
85–86

There is a typo here: op -> block.

92

Side note: why aren't we using early increment here? This can be done in separate revision but it could be dangerous/annoying to have subtly different iterator invalidation semantics for these different versions. Also, I think these semantics aren't documented on the walk methods.

We have the same issue at at least one more instance above as well.

This revision now requires changes to proceed.Feb 23 2021, 5:53 PM
bondhugula accepted this revision.Feb 23 2021, 5:56 PM

This is nice! Thanks. Just a typo I caught. I think we should fix the iteration invalidation semantics in a subsequent revision (I can do that if you want.)

This revision is now accepted and ready to land.Feb 23 2021, 5:56 PM
dcaballe updated this revision to Diff 325989.Feb 23 2021, 10:33 PM
dcaballe marked 2 inline comments as done.
dcaballe edited the summary of this revision. (Show Details)
  • Passing order as an extra parameter.
  • Fixing typo.

I'll commit this tomorrow if no further comments.

mlir/lib/IR/Visitors.cpp
85–86

removed

92

Good point! IIUC, the callback here is on the block. We invoke the callback before or after iterating over all the block's operations so the operation iteration wouldn't be invalidated even if the callback erased block's operations. It could invalidate the block iterator, though, if the block is erased... Yeah, it's a bit inconsistent. Maybe we should use early increments for all the region/block/operation iterators? I could quickly take care of it in a separate patch.

dcaballe planned changes to this revision.Feb 26 2021, 7:00 PM

Let me revisit a few things in this patch and come back.

dcaballe updated this revision to Diff 327347.Mar 1 2021, 8:59 PM

This new diff introduces important changes to the previous version.
Hopefully, @rriddle/@frgossen/@bondhugula can take another look.

  • Added new TestIRVisitors pass with basic testing for the IR visitors (I didn't feel very comfortable without it).
  • Introduced 'Prune' as a new WalkResult. It should be used, for example, to skip the walk of an operation that has been erased in pre-order. Otherwise, the walk would recurse into the erased operation (thanks for reporting the issue, @sgrechanik!).
  • Addressed @bondhugula's comment on early increment ranges. They are only needed for block and operation ranges since regions can't be directly erased (right?).
  • Updated documentation to include erasure information and walk order information. The walk order information was misleading since it only refers to nested operations/block/regions. Operations in a block are always visited in lexicographical order regardless of the nesting order. Same for all the blocks in a region or all the regions in an op.

Please, let me know if the new doc makes more sense. We can iterate on it.
Thanks!

This revision is now accepted and ready to land.Mar 1 2021, 8:59 PM
dcaballe requested review of this revision.Mar 1 2021, 8:59 PM

Can you separate out the prune functionality into a different commit?

Can you separate out the prune functionality into a different commit?

I could commit the previous version of the patch as is (already approved) and then create a new review with the new changes. Would that work better?
If I removed only the prune part, the new testing pass would fail since we need the prune functionality for pre-order walks that erase ops.

Can you separate out the prune functionality into a different commit?

I could commit the previous version of the patch as is (already approved) and then create a new review with the new changes. Would that work better?
If I removed only the prune part, the new testing pass would fail since we need the prune functionality for pre-order walks that erase ops.

Yes, I would prefer to review the pruning stuff separately.

dcaballe updated this revision to Diff 327607.Mar 2 2021, 3:20 PM

Restoring previous approved version. I'll commit it if no more comments.
I'll create a separate review for the prune changes.

rriddle accepted this revision.Mar 2 2021, 3:24 PM

LGTM

This revision is now accepted and ready to land.Mar 2 2021, 3:24 PM