[mlir][Linalg] Canonicalize decomposed ops by default.

[mravishankar]

Summary

The decomposition is meant to work hand-in-hand with
canonicalization. Add the canonicalization patterns to the pattern set
in populateDecomposeLinalgOpsPatterns.

[stellaraccident]

Makes sense to me in theory. Is there a kind of case we can test for that needs this?

[tpopp]

This creates a situation where unrelated linalg::generic ops can be canonicalized as a side effect. This has happened multiple times in linalg infrastructure and is always confusing to a user when the populated patterns are actually doing much more than they claim to do because of these global canonicalizations. I can understand the benefits of this when only applied to the newly created ops, but as is this creates confusion.

[jpienaar]

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

382 ↗

(On Diff #450225)

This seems more generic than this pass/could be useful other spots. It's also not specific to decompositions, it's specific to genericop instead (well, I'm assuming from name), and canonicalize on populate method could have instead been related to "composed ops". And in particular as this is a populate method one is probably already calling multiple populate methods on the rewrite set, and this call would just be another one.

[mravishankar]

@tpopp with pattern rewriter that is not possible. You could do it using transform dialect in this case. Here the decomposition explicitly relies on a canonicalization of linalg.generic operation that deduplicates operands and removes unsused results. That makes this transformation much more "mechanical" instead of having to reason about multiple things. So it makes sense to do the canonicalization as part of "populate" methods.
Your broader point about canonicalizations is valid though (I have hit it many times too). I broadly take two approaches to it

1. Maybe the canonicalization is not really a canonicalization. It is making an opinionated change that is not a canonicalization. These should be moved out of canonicalization. In MLIR as a whole we have a lot of things as canonicalization pattern that shouldnt be.
2. If the analysis that is being done as part of a transformation is being tripped up by canonicalization, then I'd revisit the analysis. Canonicalization is supposed to make the op/IR more easy to analyse. An analysis that works on a non-canonicalized operation should also work on a canonicalized operation.

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

382 ↗

(On Diff #450225)

Sorry, wasnt able to parse the comment here. Could you clarify?

[jpienaar]

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

382 ↗

(On Diff #450225)

Sure. The reason we use populate methods is to be directed helpers that enable composing solutions/conversions elsewhere. They are simple and just enables creating a composed pass, they all have single responsibility: populate the requested set of patterns. Without an additional param here one has a single call at populate spot that makes the behavior explicit, while here it is implicit that canonicalize does not relate to the ops corresponding to the patterns being populated or ops necessarily generated from them, but to a generic op so you could have an input without any ops affected by the patterns here than now get impacted. Which makes one then wonder: why not have the same param on all linalg populate methods? Your reasoning above would seem to apply that is desirable. But that would seem a bit redundant. Why not alternatively always add canonicalization patterns to the bag of patterns being populated at populate site? (Or could even be a specialization of the driver that additionallg uses all canonicalization patterns ...).

[jpienaar]

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

382 ↗

(On Diff #450225)

(not so say we shouldn't improve populate methods, at least run into wanting to be able to specify with a baseline benefit value and filtering out some patterns, but those are closely related to patterns being populated)

[stellaraccident]

I did add the canonicalization patterns to a downstream and have observed that they do "help" a lot of real cases I was seeing get to a fully decomposed state.

Since it seems like you primarily want to test that the canonicalization patterns "play nice" or constructively interfere when combined, why not just have the option at the test pass level and include the patterns there for testing purposes (versus plumbing that through the populate)?

[stellaraccident]

mlir/test/lib/Dialect/Linalg/TestLinalgDecomposeOps.cpp

37

If keeping this, the name is misleading: it canonicalizes everything, even if not decomposed.

[tpopp]

I did add the canonicalization patterns to a downstream and have observed that they do "help" a lot of real cases I was seeing get to a fully decomposed state.

Just to note, I fully agree that canonicalizations are beneficial in many cases and canonicalizations before/after a pattern are very useful.

So it makes sense to do the canonicalization as part of "populate" methods.

These patterns can currently be added by the caller if desired in their pipelines and if they accept the consequences of this. I acknowledge that statement also applies to this patch, but it is cluttering up the given API of a core MLIR function without also pushing for improvements to the core infrastructure. I think this change would be completely valid in a downstream project working around the core repository, but it seems less valid to create workarounds in the core repository because of shortcomings in the core infrastructure.

Here the decomposition explicitly relies on a canonicalization of linalg.generic operation that deduplicates operands and removes unsused results.

Would it be possible to have this as a separate function like populateLinalgExpectedCanonicalizations? It looks more annoying than a second parameter, but the difference in my mind is that the original function operates on some set of Ops A and doesn't operate on other Ops B, so when I see a canonicalize pattern, I expect the canonicalizations are being applied to A and not B because other behavior would be very strange (without knowing the limitations of the pattern rewriter and that this does not somehow work around it).

[mravishankar]

Stepping back a bit, I am not terribly bought into this patch. I confess, I didnt put that much thought into it, apart from the consideration this decomposition patterns works best in conjunction with the linalg::generic canonicalization pattern. It would mean every user of this pattern will also have to run canonicalization after (or add canonicalization pattern to the pattern set) to get the desired result. It made sense to me to then just add the canonicalization pattern by default. I understand that this would mean canonicalization would also run on other linalg::generic operations not created by this pattern. If that is not desriable, I am happy to drop this patch. I personally dont think that it is a problem. When a pass is run on scope, all operations within the pass scope are modifiable. It is also fair to say then its upto pass author to pull in the decomposition patterns + canonicalization pattern . I was just thinking of it as "this would be what everyone would want/should do anyway, so just make that the default behavior of the populate* method". I can step away from this direction as well ( I do understand the arguments made against including canonicalizations this way).

[stellaraccident]

I'm ok/supportive of having the option in the test pass to run with canonicalizations mixed in: it isn't a bad idea to be and to test that they work together. But yeah, I'm not thrilled about coupling that in the actual populate method.

[mravishankar]

Address comments.

[mravishankar]

mlir/test/lib/Dialect/Linalg/TestLinalgDecomposeOps.cpp

37

I didnt change the name cause for the test pass it is expected to work on one op, and decompose it + canonicalize all. If you think the name is misleading I can change it.

[tpopp]

I'm resigning as reviewer as I just wanted to voice my concerns regarding populate methods creating global effects. I wish more tests used the canonicalizer just to ease the mess that some FileCheck tests are, and I think targeted canonicalizations can make sense.

It would mean every user of this pattern will also have to run canonicalization after (or add canonicalization pattern to the pattern set) to get the desired result.

This is a common pattern that I see in MLIR. No strong feelings, but just pointing out that it's not a new situation in general.

mlir/test/Dialect/Linalg/decompose-ops.mlir

2

Is this running the pass twice now? Also, this doesn't set the canonicalize flag to true currently.

mlir/test/lib/Dialect/Linalg/TestLinalgDecomposeOps.cpp

37

No strong feelings from me either way (mentioning it as I previously mentioned a more explicit name).

mlir/test/Dialect/Linalg/decompose-ops.mlir

-// RUN: mlir-opt -test-linalg-decompose-ops -cse -split-input-file %s | FileCheck %s
-// RUN: mlir-opt -test-linalg-decompose-ops -cse -canonicalize -split-input-file %s | FileCheck %s --check-prefix=CANONICALIZECHECK
+// RUN: mlir-opt -pass-pipeline="test-linalg-decompose-ops{canonicalize-decomposed-ops=false}, cse" -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -pass-pipeline="test-linalg-decompose-ops, cse" -test-linalg-decompose-ops -cse -canonicalize -split-input-file %s | FileCheck %s --check-prefix=CANONICALIZECHECK
 
 func.func @simple_op(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>)
     -> (tensor<?x?xf32>, tensor<?x?xf32>) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
   %d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
   %init1 = linalg.init_tensor [%d1, %d0] : tensor<?x?xf32>

mlir/test/lib/Dialect/Linalg/TestLinalgDecomposeOps.cpp

struct TestLinalgDecomposeOps
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<AffineDialect, linalg::LinalgDialect>();
   }
   StringRef getArgument() const final { return "test-linalg-decompose-ops"; }
   StringRef getDescription() const final {
     return "Test Linalg decomposition patterns";
   }
 
+  Option<bool> canonicalizeDecomposedOps{
+      *this, "canonicalize-decomposed-ops",
+      llvm::cl::desc("Canonicalize the decomposed ops."),
+      llvm::cl::init(false)};
+
   void runOnOperation() override {
     MLIRContext *context = &this->getContext();
     RewritePatternSet decompositionPatterns(context);
     linalg::populateDecomposeLinalgOpsPattern(decompositionPatterns);
+    if (canonicalizeDecomposedOps) {
+      linalg::GenericOp::getCanonicalizationPatterns(decompositionPatterns,
+                                                     context);
+    }
     if (failed(applyPatternsAndFoldGreedily(
             getOperation(), std::move(decompositionPatterns)))) {
       return signalPassFailure();
     }
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestLinalgDecomposeOps() {
   PassRegistration<TestLinalgDecomposeOps>();
 }
 } // namespace test
 } // namespace mlir