This is an archive of the discontinued LLVM Phabricator instance.

[mlir] GreedyPatternRewriter: Add ancestors to worklist
ClosedPublic

Authored by springerm on Dec 19 2022, 6:33 AM.

Details

Summary

When adding an op to the worklist, also add its ancestors to the worklist. This allows for RewritePatterns to match an op a based on what is inside of the body of a.

This change fixes a problem that became apparent with vector.warp_execute_on_lane_0, but could probably be triggered with similar patterns. On a high level, the pattern extracts an op b with eligible = true from the body of an op a:

test.a {
  %0 = test.b() {eligible = true}
  yield %0
}

Afterwards:

%0 = test.b() {eligible = true}
test.a {
  yield %0
}

The pattern is an OpRewritePattern<OpA>. For some reason, test.a is not on the GreedyPatternRewriter's worklist. E.g., because no pattern could be applied and it was removed. Now, another pattern updates test.b, so that eligible is changed from false to true. The OpRewritePattern<OpA> could now be applied, but (without this revision) test.a is still not on the worklist.

Note: In the above example, an OpRewritePattern<OpB> could have been used instead of an OpRewritePattern<OpA>. With such a design, we can run into the same problem (when the eligible attr is on test.a and test.b is removed from the worklist because no patterns could be applied).

Note: This change uncovered an unrelated bug in TestSCFUtils.cpp that was triggered due to a change in the order in which ops are processed. A TODO is added to the broken code and test cases are adapted so that the bug is no longer triggered.

Depends On: D141365

Diff Detail

Event Timeline

springerm created this revision.Dec 19 2022, 6:33 AM
springerm requested review of this revision.Dec 19 2022, 6:33 AM
Herald added a project: Restricted Project. · View Herald TranscriptDec 19 2022, 6:33 AM

Note: As an alternative to this change, it may also be possible to limit what users are allowed to do inside matchAndRewrite.

springerm edited the summary of this revision. (Show Details)Dec 19 2022, 6:34 AM
springerm edited the summary of this revision. (Show Details)

Note: As an alternative to this change, it may also be possible to limit what users are allowed to do inside matchAndRewrite.

FWIW, we already have such limitations, e.g., don't bypass the rewriter. But since they are only English, many users just ignore them. I don't see an easy compiler-assisted way to enforce these: once the user gets ahold of an Operation *, they can do arbitrary mutation with the low-level API.

rriddle added inline comments.Dec 20 2022, 10:33 PM
mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
567–568

This isn't sound. You can easily end up adding parent operations that shouldn't be added, e.g. if you're in a function pass you shouldn't be adding the parent function/module/etc. This needs to gate on the upper limit of what can actually be considered.

springerm planned changes to this revision.Dec 21 2022, 7:18 AM
springerm updated this revision to Diff 485947.Jan 3 2023, 2:50 AM

address comments

springerm marked an inline comment as done.Jan 9 2023, 1:14 AM

@rriddle Can you take another look at this?

Thanks for the ping, just coming out of vacation today. Looks reasonable to me, but curious about the cost of the check.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
313–322

Do you have any idea of the relative cost of this check? The pattern rewriter is fairly hot, just wondering if this has a visible impact on compile time.

mehdi_amini added inline comments.Jan 9 2023, 9:43 AM
mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
106

Can you remove this and make it an argument of the addToWorklist method? Seems to me this is only every local to the simplify method and adding "global" mutable state like does not seem ideal to me.
(also should be an ArrayRef I believe)

mehdi_amini added inline comments.Jan 9 2023, 9:47 AM
mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
106

Actually I'm incorrect: this is called from many places, it may makes sense to have it here but please document the invariant a bit more:

/// This is set at the beginning of `simplify()` to the current scope the rewriter operates on.

I'm still bothered by this in the context of a class that:

  1. has many public methods and simplify() isn't clearly the only entry point.
  2. addToWorklist is a virtual method but now the default implementation is using a private field member here.

Seems like too many foot guns involved here...

mehdi_amini added inline comments.Jan 9 2023, 9:52 AM
mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
313–322

We could also avoid repeating the tree-traversal of getParentOp hidden in findAncestorOpInRegion by having scope being a set and then do a single traversal but for each parent region querying the set.

That said it wouldn't be a win with a single region which may be the common case? (the code could handle it by branching on whether there is a single element in the set)

Also a nit: but Region::isAncestor(Region *other) seems a more suitable API than findAncestorOpInRegion here.

springerm updated this revision to Diff 487729.Jan 10 2023, 2:54 AM
springerm edited the summary of this revision. (Show Details)

address comments

springerm edited the summary of this revision. (Show Details)Jan 10 2023, 2:54 AM
springerm edited the summary of this revision. (Show Details)
springerm edited the summary of this revision. (Show Details)
springerm marked an inline comment as done.Jan 10 2023, 3:01 AM

Added a dependent change. It is needed so that I could write the test case.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
106

I could prepare another cleanup that sets all fields already in the constructor. simplify would no longer take any parameters. Some of the fields can be const.

Then allow simplify to be called only once. This is how we use this class anyway. We create an instance, call simplify, then it is no longer used. Does that sound good?

313–322

I think it pays off even with just one region because we have now only a single while loop that traverses the ancestors. Before, it was two nested while loops. As for the loop structure, this is as efficient as it gets.

mehdi_amini accepted this revision.Jan 10 2023, 3:07 AM

LGTM, but let's wait for River to have confirmation.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
324–328

Maybe slightly simplify and reduces the bracing.

This revision is now accepted and ready to land.Jan 10 2023, 3:07 AM
This revision was landed with ongoing or failed builds.Jan 13 2023, 1:51 AM
This revision was automatically updated to reflect the committed changes.
springerm marked an inline comment as done.

LGTM, but let's wait for River to have confirmation.

Did you get confirmation from River?

LGTM, but let's wait for River to have confirmation.

Did you get confirmation from River?

There was no further activity for multiple days, so I assume this is fine by him.

That isn't really what I intended with my conditional approval, so please get explicit approval today from River ASAP.

That isn't really what I intended with my conditional approval, so please get explicit approval today from River ASAP.

@rriddle Is this good or should I revert?

Looks reasonable to me! (Sorry I was OOO for a bit)