Page MenuHomePhabricator

[GreedPatternRewriter] Preprocess constants while building worklist when not processing top down
ClosedPublic

Authored by rriddle on Mar 29 2022, 4:28 PM.

Details

Summary

This avoids accidentally reversing the order of constants during successive
application, e.g. when running the canonicalizer. This helps reduce the number
of iterations, and also avoids unnecessary changes to input IR.

Fixes #51892

Diff Detail

Event Timeline

rriddle created this revision.Mar 29 2022, 4:28 PM
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added projects: Restricted Project, Restricted Project. · View Herald Transcript
rriddle requested review of this revision.Mar 29 2022, 4:28 PM
rriddle retitled this revision from [GreedPatternRewriteDriver] Preprocess constants when building the worklist to [GreedPatternRewrite] Preprocess constants while building worklist when not processing top down.Mar 29 2022, 4:30 PM
rriddle added reviewers: jpienaar, mehdi_amini.
jpienaar accepted this revision.Mar 29 2022, 5:01 PM

Nice, this was a pain point for folks diff'ing across passes, thanks!

mlir/include/mlir/Transforms/FoldUtils.h
51

opportunity

56

Success and failure is a bit dubious for me here. Returning a bool indicating folded & optionally hoisted feels more descriptive even.

mlir/lib/Transforms/Utils/FoldUtils.cpp
82

Could we make referencedDialect.count into a lamba helper such as alreadyFolded to be more self-documenting?

82

So here we could bubble up above not-known folded ops?

mlir/test/Transforms/test-operation-folder.mlir
27

Could we run -test-patterns 2x above too? (this is probably sufficient, but running 2x is probably more explicit here)

This revision is now accepted and ready to land.Mar 29 2022, 5:01 PM
okkwon added a subscriber: okkwon.Mar 30 2022, 5:12 PM
rriddle updated this revision to Diff 419531.Mar 31 2022, 11:55 AM
rriddle retitled this revision from [GreedPatternRewrite] Preprocess constants while building worklist when not processing top down to [GreedPatternRewriter] Preprocess constants while building worklist when not processing top down.
rriddle marked 5 inline comments as done.
This revision was landed with ongoing or failed builds.Mar 31 2022, 12:12 PM
This revision was automatically updated to reflect the committed changes.

This change is causing some behaviour which I am struggling to describe. I'm trying to make a reproducer. No need to revert (I don't think).

So there are at least 2 significant changes to behaviour

  1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.
  2. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

So there are at least 2 significant changes to behaviour

  1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.

Was this kind of awkwardly relying on the fact that we always did 2 iterations (given that we refolded existing constants)?

  1. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

The situations where that happens feels kind of wrong. Folding a constant into a different constant feels weird.

mlir/lib/Transforms/Utils/FoldUtils.cpp
82

Much cleaner, thanks.

82

Yep, just ensures that unique'd constants are pooled at the top. Should never cause issues given that constants can't have operands.

So there are at least 2 significant changes to behaviour

  1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.

Was this kind of awkwardly relying on the fact that we always did 2 iterations (given that we refolded existing constants)?

  1. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

The situations where that happens feels kind of wrong. Folding a constant into a different constant feels weird.

To expound upon this further, it feels like an implicit canonicalization on the foo.my_const.

So there are at least 2 significant changes to behaviour

  1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.

Was this kind of awkwardly relying on the fact that we always did 2 iterations (given that we refolded existing constants)?

And one of a or b were constant? Then isn't this expected to work? E.g., the constant operand is flipped before pattern matching is attempted?

  1. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

The situations where that happens feels kind of wrong. Folding a constant into a different constant feels weird.

To expound upon this further, it feels like an implicit canonicalization on the foo.my_const.

Yeah this is unfortunate, but it meant that previously this was seen as canonical behavior. Not ideal I agree. @Mogball does this just affect unit test rather than semantics/end-to-end behavior? (The new behavior probably results in less changes compared to old - in particular there was the case of discardable attributes getting dropped from constants due to this).

So there are at least 2 significant changes to behaviour

  1. The folder actually sorts commutative ops, but this doesn't trigger the greedy rewriter to run another iteration because tryToFold might still return failure. If -my-pass only has one pattern commutative(a, b) -> foo but the input IR is commutative(b, a), the pattern doesn't match because the folder will re-order the operands but the rewriter will exit.

Was this kind of awkwardly relying on the fact that we always did 2 iterations (given that we refolded existing constants)?

And one of a or b were constant? Then isn't this expected to work? E.g., the constant operand is flipped before pattern matching is attempted?

Add few more contexts here. My understanding is, we changed the timing of adding a constant to the folder,
In the past, we add a constant when it calls tryToFold but now we add a constant when it's pushed to the worklist(before calling tryToFold)

The case we are seeing is a pattern like
def Pat<(FooOp BarOp, ConstantOp), ...>
The ir is (notice the operand order)

%cst = ...
%bar = ...
%foo = Foo(%cst, %bar)

In the past, it will have two runs on the ir because it will add %cst to the folder and veiw it as a success folding. besides, In the first run, it also reorders the operands of FooOp. As a result, in the second round, the pattern can be matched.

And now, the %cst has been added to the folder. When it calls tryToFold, it just returns failure(). So there's no second round (even the operands of FooOp has been reordered). And when we check the IR that failed to match, it shows

%cst = ...
%bar = ...
%foo = Foo(%bar, %cst)

Which make us think why it wasn't matched even the form looks correct.

  1. Constants that are already in the IR aren't getting "folded". I.e. foo.my_const could be "folded" into arith.constant.

The situations where that happens feels kind of wrong. Folding a constant into a different constant feels weird.

To expound upon this further, it feels like an implicit canonicalization on the foo.my_const.

Yeah this is unfortunate, but it meant that previously this was seen as canonical behavior. Not ideal I agree. @Mogball does this just affect unit test rather than semantics/end-to-end behavior? (The new behavior probably results in less changes compared to old - in particular there was the case of discardable attributes getting dropped from constants due to this).

Add few more contexts here. My understanding is, we changed the timing of adding a constant to the folder,
In the past, we add a constant when it calls tryToFold but now we add a constant when it's pushed to the worklist(before calling tryToFold)

The case we are seeing is a pattern like
def Pat<(FooOp BarOp, ConstantOp), ...>
The ir is (notice the operand order)

%cst = ...
%bar = ...
%foo = Foo(%cst, %bar)

In the past, it will have two runs on the ir because it will add %cst to the folder and veiw it as a success folding. besides, In the first run, it also reorders the operands of FooOp. As a result, in the second round, the pattern can be matched.

And now, the %cst has been added to the folder. When it calls tryToFold, it just returns failure(). So there's no second round (even the operands of FooOp has been reordered). And when we check the IR that failed to match, it shows

%cst = ...
%bar = ...
%foo = Foo(%bar, %cst)

Which make us think why it wasn't matched even the form looks correct.

To me this is a breakage of the API contract of applyPatternsAndFoldGreedily: it does not seem right that applyPatternsAndFoldGreedily "converges" and returns while a pattern would still apply.

Add few more contexts here. My understanding is, we changed the timing of adding a constant to the folder,
In the past, we add a constant when it calls tryToFold but now we add a constant when it's pushed to the worklist(before calling tryToFold)

The case we are seeing is a pattern like
def Pat<(FooOp BarOp, ConstantOp), ...>
The ir is (notice the operand order)

%cst = ...
%bar = ...
%foo = Foo(%cst, %bar)

In the past, it will have two runs on the ir because it will add %cst to the folder and veiw it as a success folding. besides, In the first run, it also reorders the operands of FooOp. As a result, in the second round, the pattern can be matched.

And now, the %cst has been added to the folder. When it calls tryToFold, it just returns failure(). So there's no second round (even the operands of FooOp has been reordered). And when we check the IR that failed to match, it shows

%cst = ...
%bar = ...
%foo = Foo(%bar, %cst)

Which make us think why it wasn't matched even the form looks correct.

To me this is a breakage of the API contract of applyPatternsAndFoldGreedily: it does not seem right that applyPatternsAndFoldGreedily "converges" and returns while a pattern would still apply.

Agreed here. It seems to me that this has been an underlying issue masked by the way we were processing constants (i.e. unnecessarily refolding).

Sent out https://reviews.llvm.org/D122870 for a repro.

Talked to River offline: seems like a good reason to temporarily revert, he'll re-land with a fix.