Page MenuHomePhabricator

[MLIR] Introduce a new rewrite driver to simplify supplied list of ops
ClosedPublic

Authored by bondhugula on Jul 18 2021, 3:18 AM.

Details

Summary

Introduce a new rewrite driver (MultiOpPatternRewriteDriver) to rewrite
a supplied list of ops and other ops. Provide a knob to restrict
rewrites strictly to those ops or also to affected ops (but still not to
completely related ops).

Discourse post
https://llvm.discourse.group/t/introducing-a-new-rewrite-driver-to-work-on-a-specified-list-of-ops/3899

This rewrite driver is commonly needed to run any simplification and
cleanup at the end of a transforms pass or transforms utility in a way
that only simplifies relevant IR. This makes it easy to write test cases
while not performing unrelated whole IR simplification that may
invalidate other state at the caller.

The introduced utility provides more freedom to developers of transforms
and transform utilities to perform focussed and local simplification. In
several cases, it provides greater efficiency as well as more
simplification when compared to repeatedly calling
applyOpPatternsAndFold; in other cases, it avoids the need to
undesirably call applyPatternsAndFoldGreedily to do unrelated
simplification in a FuncOp.

Update a few transformations that were earlier using
applyOpPatternsAndFold (SimplifyAffineStructures,
affineDataCopyGenerate, a linalg transform).

TODO

  • Doc will be updated close to completion of review of this revision.
  • OpPatternRewriteDriver can be removed as it's a special case of MultiOpPatternRewriteDriver, i.e., both can be merged.

Diff Detail

Event Timeline

bondhugula created this revision.Jul 18 2021, 3:18 AM
bondhugula requested review of this revision.Jul 18 2021, 3:18 AM

Minor tweaks. NFC.

bondhugula edited the summary of this revision. (Show Details)Jul 18 2021, 3:25 AM

NFC. Improve comment.

Improve comment. NFC.

mehdi_amini added inline comments.Jul 19 2021, 6:17 PM
mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
399–401

Is this exact? Wouldn't the simplification of the operands added to the wordlist recursively triggers adding their own operands after simplification?

436

Nit: this is misleading as op is a pointer, Operation * ?

444

Can you tell me a bit more about when is this happening?

452

It should be removed of workOps here I think, otherwise another pattern can create ops and reuse the pointer that is freed now and we'd think it is part of workOps.
(in general keeping dangling pointers in set can be annoying to debug because it leads to bugs that are not deterministic)

467

Nit: I'm unsure that the autos here improves readability

468

The API supports contains now I believe

475

Can you match the name of the API? I had to look to see that it is processGeneratedConstants which is more explicit about what will happen here: it isn't just any op that we expect here but only constants (which makes sense, we don't expect the folder to create any other kind of ops).

489

Again: shouldn't it be removed from the workOps?

Also does inPlaceUpdate really implies that the op is removed? I would think that folding some attributes in-place would also result in this code path.

491

Nit: spurious braces

bondhugula added inline comments.Jul 19 2021, 8:51 PM
mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
489

@mehdi_amini It's being removed from workOps in notifyOperationRemoved :-) It has to be done there otherwise erasure happening in canonicalization patterns would miss updating workOps. Likewise for your comment above.

bondhugula marked 9 inline comments as done.Jul 19 2021, 10:56 PM
bondhugula added inline comments.
mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
399–401

Yes, this is exact. When strict is on, we prevent "outside" ops from getting into the worklist whether operand defs or result users.

436

Thanks. I missed it during copy and paste.

444

This code is common to all the rewrite drivers. When eraseOp is called, the ensuing hooks mark the item in the worklist null.

452

It's being done in notifyOperationRemoved cleanly! :-)

467

Thanks. Removed auto. (This part was was copied over from the other drivers.)

468

Thanks.

475

Sure, thanks.

491

I let the brace because there is a comment. Per LLVM style, this is recommended. (https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements)

bondhugula marked 8 inline comments as done.

Address review comments.

count(..) > 0 -> contains

herhut accepted this revision.Jul 20 2021, 1:36 AM
herhut added a subscriber: herhut.

Cool, thanks for adding this. Please also wait for @mehdi_amini to agree before landing.

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
81

nit: as well as

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

I find workOps not very descriptive. It is a working set, so I kinda get that name workOps but when reading the code further below I find it confusing with also a worklist.

Maybe strictModeSet or strictModeFilter would be more descriptive?

413

nit: as well as

422

nit: remove once

423

nit: is to keep

This revision is now accepted and ready to land.Jul 20 2021, 1:36 AM
mehdi_amini accepted this revision.Jul 20 2021, 11:35 AM
mehdi_amini added inline comments.
mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp
90

(Nit: do you know that SmallVector size is optional now?)

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

Good point on the naming, I would throw in the pot the suggestion for strictModeFilteredOps :)

452

Ah, sweet! :)

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp
555–556

Comment can be updated I think.

bondhugula marked 9 inline comments as done.

Address review comments from @herhut and @mehdi_amini.

mlir/lib/Dialect/Affine/Transforms/SimplifyAffineStructures.cpp
90

Yes, but a force of habit here! Dropped it.

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

Makes sense. I was also thinking of strictModeFilteredOps after @herhut suggestion.

lattner resigned from this revision.Jul 21 2021, 7:52 AM
This revision was landed with ongoing or failed builds.Jul 21 2021, 7:56 AM
This revision was automatically updated to reflect the committed changes.

Thanks for the comments @mehdi_amini and @herhut. I'll have a follow-up to add the documentation for this.