RFC : Update to "General Design" section of Operation Canonicalizations in MLIR

MLIR
1

Over the past few months there have been a lot of discussions on PRs that attempt to add canonicalizations. To avoid having long discussions on PRs, and to provide better guidance to new contributors, it might be time to update the canonicalization section of MLIR documentation. While the existing documentation provides some good basics for how to use canonicalizations, drawing inspiration from Dan Gohman’s article on canonicalizations, it does not provide more emphasis on the “Do No Harm” principle specified in that article. In my view most of the discussions have been centered around not having spelled this out clearly.

Proposed basic guidelines for something to be a canonicalization

Canonicalizations in MLIR should be designed in a way that they could be applied repeatedly, and at any level of the compilation stack. For this to be true, canonicalizations should be geared towards converting the input program into a state that is almost always beneficial (i.e. do no harm). The post already talks about cases where things are ambiguous (I find the point about redundancy elimination also extremely important and something I have seen have huge impact in practice). So some ambiguity here is unavoidable, but in my opinion the following thumb rules help

  1. Transformations need to have a relatively high bar for being a canonicalization. Conversely, there should be a relatively low bar to have transformations moved out of canonicalizations.
  2. Transformation that are dropping semantic information captured by operations should not be part of canonicalizations.
  3. Transformations that are just changing program representation, without necessarily making the program simpler should not be a canonicalization.

These are just a few to start with. I am happy to adapt/evolve them based on discussions here, and will update the General Design section of the Operation Canonicalization page accordingly.

Following are examples from a past discussions where the above thumb rule could be used to reach a decision.

tensor.expand_shape/collapse_shape vs tensor.extract_slice/insert_slice

This discussion came up recently on a PR. The change was attempting to add a tensor.insert_slice of the form

%0 = tensor.insert_slice %slice into
     %x[0, 0, 0, 0, 0][1, 1, 1, 16, 32][1, 1, 1, 1, 1] :
     tensor<16x32xf32> into tensor<1x1x1x16x32xf32>

to

 %0 = tensor.expand_shape %slice[[0,1,2,3], [4]] :
          tensor<16x32xf32> into tensor<1x1x1x16x32xf32>

These are equivalent representations. There is nothing to say which is better. It is equally valid for an analysis to uses tensor.extract_slice/insert_slice in a load-bearing manner (for example One Shot Bufferization in MLIR), as opposed to using tensor.collapse_shape/expand_shape. Moving from one form to another would break the subsequent program analysis; or force analysis to account for these two being essentially equivalent. The latter is essentially “doing harm” since the analysis would have to essentially undo the transformation applied as canonicalization. Furthermore one could add a transformation that converts tensor.expand_shape to tensor.insert_slice also as a canonicalization resulting in an infinite loop during canonicalization (without explicit break out). The fundamental issue here is that there is nothing that makes one representation “more canonical” than the other and should not something that canonicalizations do.

vector.transpose to vector.shape_cast

This commit added a pattern to convert vector.transpose of the form

vector.transpose %0, [1, 0] : vector<nx1x<eltty>> to vector<1xnx<elty>>

to a

vector.shape_cast %0 : vector<nx1x<eltty>> to vector<1xnx<elty>>

The motivation for this change was when lowered to LLVM, the extra-unit dimensions dont really matter. From LLVM backend’s perspective the source type of the vector.transpose and the destination type are the same, and hence this vector.transpose is a no-op and converting this to a vector.shape_cast essentially captures this “no-op” behavior. The problem with this though is that it violates the “canonicalizations shouldnt drop semantic information captured by an operation”. The vector.transpose version of the operation captures more semantic information about how the dimensions are changed, than the vector.shape_cast operation. So adding this as a canonicalization would result in this semantic information being dropped much earlier than it should be. It is true that such transposes should be a no-op for LLVM-based lowering. Such a transformation should then be run as part of, or as a step immediately preceding, the conversion to LLVM dialect. To make the distinction more concrete, consider

vector.transpose %0, [1, 0, 2] : vector<nx1x1x<eltty>> to vector<1xnx1x<elty>>

If this was converted to a vector.shape_cast this would be

vector.shape_cast %0 : vector<nx1x1x<eltty>> to vector<1xnx1x<elty>>

It would be unclear if this vector.shape_cast came from vector.transpose as written above or from this one

vector.transpose %0, [2, 0, 1] : vector<nx1x1x<eltty>> to vector<1xnx1x<elty>>

This ambiguity is a result of semantic information captured by the vector.transpose operation being dropped. Hence the change from vector.transpose → vector.shape_cast should not qualify as a canonicalization.

Good use of canonicalization

To highlight a good use of canonicalization, propagating static shapes in the program is a good use of canonicalization. This makes it easier to analyse programs and leads to a more compact representation.

Waiting to hear back from the community on their take on this.

2

Speaking specifically about insert_slice vs expand_shape case, I think, actually, expand_shape is strictly better than insert_slice because it drops %x use, opening more opportunities for other canonicalizations (DCE). And in real code this arg can come from actual chain of ops which will be dead otherwise instead of single init_tensor in case of one-shot. So you are blocking users optimizations because of questionable decisions of single part of (potentially unrelated) infrastructure.

(also, I’m not a big fan of one-shot bufferization design in general and I also think it was heavily misadvertised, but it’s a rant for another day)

3

On this point specifically, I think we need to stop guiding ourselves over a document written on top of a blog post, with some opinionated ideas on what canonicalization is, and start discussing what canonicalization really means for us, for MLIR and for our projects using it.

To me, there are two world views in MLIR, each require a different form of canonicalizaiton:

  1. Downstream compilers that use MLIR for building new dialects, which have their own canonical forms and need custom canonicalization mechanisms, not many provided by upstream MLIR.
  2. Upstream passes and downstream compilers that use MLIR for their dialects, interfaces and transforms, and that rely on an upstream canonical form for those dialects in order to avoid combinatorial explosion. MLIR does not provide a good solution here, either, with confusing statements about what canonicalization means.

The former needs fine-grained control of how to create custom canonical forms, the latter needs to use that framework and agree on a set of rules so that upstream passes/transforms can interoperate with each other and with downstream passes/transforms.

So, in my opinion, we need to look at both, but separate the discussion between the infrastructure and the agreed-upon upstream canonical form on the upstream dialects.

Later, we also need to make sure we discuss not only operation canonical forms but sub-graph / function / module canonical forms when we want to match more than just one op. But this is for another RFC.

4

Thanks for starting the discussion.

I cannot follow this: there is no way in my opinion to define what “do no harm” means. There is almost always a path or a target that will prefer one form or the other. This is the reason why Dan’s post does not emphasize this, and on the contrary I believe it explains why it isn’t a good idea to emphasize this in the “Yeah so what about x * 2…” section: it can be seen “harmful” to turn a fast addition into a slow multiplication, but canonicalization isn’t about this! That’s missing the whole point.

“do no harm” requires cost model and complex analysis which are beyond what you can expect from canonicalization. However the “do not drop information” and “can easily recover” part of the principle should already be enough! (because if you don’t lose information, how can you do harm since an optimization can recover?).

This is overly ambiguous and not very principled IMO.

I don’t understand this: there is little definition of “simpler” and that is missing the point of canonicalization entirely: the subsequent analysis and transformations should have an easier time finding patterns not because the program is simpler but because it is in a pre-defined form!!
According to this principle we should stop pushing constants to the right of commutative operations for example…

The fundamental of canonicalization is there: at a given abstraction level, two semantically equivalent program should be transformed in the same form as much as possible. If you don’t subscribe to this, then why use canonicalization at all? Especially: it’s not an “optimization”.

Here is the basic set of principles I set for canonicalization:

  1. Cannot lose information. The information should be recoverable without “heroics” or “raising” by subsequent analysis and transformations: we want these to simply pattern match the canonical form.
    As such canonicalization operates at one level of abstraction (that does not mean “inside a dialect”, it’s unfortunately ambiguous as some dialect sometimes span two level of abstractions, or one level of abstraction spans multiple dialects…).
  2. Should be fast enough: canonicalization is iterative and especially matching can’t be overly complex.

Do you mean “break” in the abstract absolute sense, or break the “current implementation of some analysis”? The former is a good argument against a canonicalization, the latter is just the usual process of landing changes in the right order (any analysis should be updated to match the canonical form…).

Actually without canonicalization, you are forcing the analysis to account for these to be equivalent. This is exactly what canonicalization solves: the optimizations does not need to know and can match only one form.

I don’t quite understand the argument here: for every single canonicalization you could say that “someone could add an infinite loop”.

This is the case with almost all canonicalization, back to Dan’s article which touches on all the ambiguity.

Right and if I remember correctly someone (@sogartar?) argument this point convincingly, because the “don’t drop semantics” part is a fundamental principle to follow.

I agree in principle, but this is a difficult one in practice because most of the time that requires adding a casts back to the dynamic type. This is because you can’t just update the result of an op without also updating all the users, otherwise you’ll break a lot of verifiers.

For example this discussion RFC: remove arith/math ops on tensors - #63 by kuhar have a consensus that people didn’t want to support arith.extsi %3 : tensor<12xi16> to tensor<?xi64>, however that means that if the original form is:
arith.extsi %3 : tensor<?xi16> to tensor<?xi64>, then inferring the shape on the operation that produces %3 can’t be a local decision otherwise it would produce the invalid arith.extsi above.

It can still work but all the “shape inference patterns” need to look through the newly inserted cast, and there can be casts left behind which then isn’t clear it is desirable? (it can be argued that static->dynamic conversion in the type system is fine as separate cast).

5

It must have been someone else.

6

I feel like the latter is the main issue here. I see these discussions mostly when stuff is about to break. I think many people are pragmatic and don’t care as long they don’t have to spend any time on the next LLVM integrate.

A first step could be to provide fine-grained access to all canonicalization patterns through populate...Patterns functions. (More fine-grained than Op::getCanonicalizationPatterns.) Maybe even some kind of pattern registry. Then projects that want to have a lower integration burden can selectively opt into only the canonicalizations that they need. (And not use the MLIR canonicalizer pass.)

7

You may be onto something!

Maybe it isn’t well-known, but the canonicalizer pass exposes filters ions, options to be able to opt in or opt out of some patterns individually by name already.
A very conservative downstream users could just opt in the patterns they want, or when broken by a change upstream opt out.

8

side note here, those patterns names were initially added for the debugging only purposes. They are generated from the pattern class name, using function pointer pattern API will result in an empty name and there may be conflicts as multiple unrelated patterns from different places may have the same name. IMO, this system is not very robust for the production use.

9

On this point specifically, I think we need to stop guiding ourselves over a document written on top of a blog post, with some opinionated ideas on what canonicalization is, and start discussing what canonicalization really means for us, for MLIR and for our projects using it.

+100. I think it would be very valuable to build this with an internally maintained set of guidance rather than referring to external things that are “right-ish”.

The practical side of canonicalizations is that are run frequently, and intermix with other canonicalizations, and are optional, I this we should include guidance that at least includes:

  1. Because canonicalize are iterated, you need a lattice like approach where you don’t get cycles in canonicalizations or else you get “infinite” compile time. You need to define a “more canonical form” and things need to agree on that.
  2. Canonicalizations should not be required for correctness. If they are, then it’s a lowering, and canonicalize hooks are the wrong place to do it.
  3. You shouldn’t do something expensive that is O(n) in canonicalize. Doing so will be a significant problem for compile time, because they are run often and are iterated. Such things are more suitable for a dedicated pass that is run in a more controlled way.
  4. Canonicalize isn’t a great place to put things with complicated cost models, because the def-use graph changes a lot and the iterative nature means that the IR is changing a lot. When combined with #2, this is one of the reasons that “tensor constant folding” is a bad thing to do in canonicalize/fold hooks.

Going beyond these high level rules, I think the “do no harm” rule is right minded, but I’m not sure how to phrase that. I also think that canonicalizations are massively over-used in the standard dialects.

In any case, it would be wonderful to propose and iterate on some specific guidelines, and show some antipatterns to provide guidance for the community. The fact that this “tribal wisdom” isn’t written down makes is some less-MLIR-experienced often struggle, because doing the wrong thing is easy and appealing and you run into problems late when writing tons of code.

-Chris

10

I don’t know how to phrase the “do no harm” guideline in a more useful way, but it does resonate with me. Personally, I think the need to state it is more of a symptom of exactly what you are pointing out: canonicalizations have been pretty widely abused in the standard dialects to the point that without an overriding design having been followed, sometimes the best feedback that a reviewer can provide on the spot is that a change may be harmful to the situation. So I interpret the intent here basically as a summary of what everyone is saying: the current state isn’t very good, we need a higher bar of some kind, and that starts by taking more of a default skeptical view of canonicalizations that are being proposed as one offs without really thinking through whole state.

I see a lot of canonicalizations being added to highly generic dialects that read to me as part of lowering pipelines that would be more logically categorized as things you do to prepare the IR for a certain kind of transformation. Often times this gets baked in because the first kind of lowering that makes its way in tree just adds a fair number of the patterns it needs as canonicalizations without really considering the lattice of canical forms.

In any case, I would +100 the diagnoses in this thread if I could but don’t have a lot specific to add that hasn’t already been expressed.

The one thing I would say is that in my experience, the further “up” you go towards generic abstractions, the higher the likelihood that there may not be as many useful canonical forms that can be expressed easily. I think this is ok, and if the choice is having very few canonicalizations for such things vs badly/partially defining a canonical form that wouldn’t otherwise exist “in nature” – I would bias towards fewer and being conservative in adding them. Some of these things are modeling concepts quite a bit more complicated than traditional scalar instruction sets, and it may just be prohibitively hard or non useful to formally define things to pass the bar. For most of those cases I have seen, I would vastly prefer to define various libraries of simplifications as a form of lowering that you do when the situation calls for it.

11

Agree, a major challenge in the MLIR community is the desire to prematurely abstract things. MLIR core is being able to solve for many levels of abstraction, and this leads to the desire by some to build extremely general infra, passes, etc etc.

This can be great for some things (e.g. canonicalization which is a universal concept across all levels, at least when done right) but it leads to a watering down and conflation of concerns as you say.

-Chris

12

I’d like to really see examples to accompany what seems to be more a fuzzy feeling about the state of canonicalization, and also quantification (the existence of a few bad patterns does not invalidate the rest of the code base!).

I don’t share this impression: while imperfect and lacking formalism and documentation, it does not seem to be that canonicalization are in such a bad state as folks are making it seem.
In particular I haven’t seen much of the discussions in the past being showing issues of “lowering” or “codegen prepare” in the canonicalization pattern so far.

It is quite the contrary actually in my experience: recurring complains seem to come instead from folks who built their lowering pipeline relying on a specific canonical form for correctness, while their expectation is not documented and the canonicalizer is far from being complete. They are then being broken when upstream improves the canonicalized: but the root of the problem here is the lack of specialization and as you say “things you do to prepare the IR for a certain kind of transformation”. Because people don’t structure their pipeline to do that they built implicit assumptions on the IR and get broken when upstream improves.
This complaints often comes under the form against a given canonicalization pattern: “but then I can’t run canonicalization right before my lowering…” : this is another illustration of missing “things you do to prepare the IR for a certain kind of transformation” before their transformation (i.e. it is expected that they need a “codegen prepare” phase and can’t just run canonicalization afterward!).

That is, they break this rule:

This isn’t an upstream canonicalization problem though, this is just that documenting invariants on IR and building a pipeline that is “complete” (as in: take all possible input IR for a dialect) is difficult (most people will take code sample and examples instead of “working through the spec” and handling a dialect in its entirety); this leads to fragile systems and I’m not sure there exist a solution for this.

13

In my interpretation of that concept, there are two different things:
While I agree canonicalization should not be required for correctness (of the IR semantics, verification, etc), it is often required for a pass to work at all (pattern matchers).

So far, I have not looked at canonical forms as a correctness thing, these things are mostly orthogonal. My intention in looking at canonicalization is not to find correctness, but to find “amenable forms for further transformation”, and that depends on the dialects we use, which interfaces they implement and what passes we’re running on them.

In that sense, there’s no true canonical form. There’s a bunch of different canonical forms for different objectives.

My argument to this discussion before the RFC was a pragmatic approach. First, we define what is the purpose of canonicalization (there may be more than one), then we define what’s the infrastructure we need to build to get there (@clattner has excellent points up-thread), then we adapt our current dialects/interfaces/passes to that, both upstream and downstream.

Same here. I think Chris’ point go towards that direction and we can start with that. We can also add less hand-wavy statements that are still “guidelines” and not concrete points, for example: “It should not lead to an ambiguous state”, (which os generally ill defined, but each op can have its own set). This can lead to cycles and complex run times.

If we come up with enough general rules, the intent to “do no harm” will be encoded in the rules and not need to be explicitly stated, which would remove the imprecision semantics and cultural load of the word “harm”.

14

I’m sorry but I can’t agree with this perspective and I think you are only seeing one side of a debate where the upstream situation has become so intolerable that most of the people doing work in the area have stopped engaging.

The upstream “standard tensor” pipelines (don’t know what else to call them), in general, are the offenders in this area, and because the bar is low/wrong in terms of supporting specific lowering paths with proper preparation and “completeness” vs the current state of being too abstracted, it kind of infects everything.

Concretely (although I’m not going to take the time to present a full argument):

  • bufferizarion is actually a suite of algorithms that should each be optional. But the specifics of this one are deeply embedded in both the canonical forms of the adjacent dialects and fragile to it (reference how many times discussions end with “bufferizarion requires it this way”).

  • The vector pipelines are similarly intertwined into a higher order thing that really should be multiple concerns and mid level optimizers tied together with a common vocabulary. The current canonicalizations are a mismatch between actual canonical forms and things that are more a semi common set of non optional prep steps for some set of things that were important to someone at some point. It’s next to impossible to tell the difference at this point vs doing what Mahesh is trying to do by straining things into specific pipelines.

In both situations, the over generalization and lack of specificity upstream has bled over into many things looking like canonicalizations and being accepted as such instead of being sunk into specializations that are themselves complete and coherent pipelines. This debt then infects a lot and makes it hard to see the soup for the nuts when looking at things patch by patch. We don’t even know to what extent the situation is salvageable: we’re investing significant time in a try to straighten some of it (this thread itself is a byproduct of making that push to try to clean things up), but may end up more in the territory of resolving it by fork of key pieces if we can’t force de-abstraction into the way that upstream operates (not even saying that this becomes a downstream fork but it may need to resolve by producing a V2 of some of this stuff vs in situ fixes – we’re just not sure yet).

In any case, it sounds like we are agreed on the root design flaw.

Where we differ is in the analysis of where the bitrot originates and the likely depth of the cuts needed to resolve it – and whether upstream got it right more than setting up sub-optimal implementation patterns that got copied and mitigated against.

While that is a topic for another day, canonicalization is a frequent battleground because of many of the things discussed here. Better guardrails and a higher bar there will at least help bias the overall design in a positive way and may provide some breathing room to actually fix some of it.

I think we should start by using Chris’s guidelines as a basis and see if we can encode some of the “do no harm aspects” as do’s and dont’s by way of example.

15

I think these two are related. Even where losing includes change a local property to one requiring a global analysis. Which feeds into the “amenable for further transformations” even. I was trying to think if “losing information” has a phase related component …

I think harm captures it to some extent. Its about introducing additional cost and engineer/compiler analysis pain. E.g., some of the cultural load may actually be appropriate :slight_smile:

From my mathematical glossary part this is scratchy :slight_smile: (even
Canonical form - Wikipedia ). Saying its canonical for a phase or pipeline or something does make sense, although then we get into arguments about what are the phases :wink: BUT this would then say from an infrastructure point of view we should potentially split the concepts. Currently an op has a canonical form. Its independent of “time and space”. And we made it very convenient to register them. If we go this route we would be saying an op has a canonical form either in general, per pass or per phase. Now the latter the op knows nothing about … Perhaps I’m now exactly going down the route of prematurely abstracting and these are just regular passes with a name that folks interleave + maybe some convenience wrapper.

This seems to be an issue with bufferization rather than canonicalization. I think folks have strong opinions about the change to the one shot approach (which I believe actually resulted in these things being less optional and more required). And supports Mehdi’s points of a missing prepare pass/patterns and that these shouldn’t be canonicalizations (so you two are in agreement, and this is a possible example of guard rail) … well unless they make it more amenable for further transformations. E.g., what bufferization needs and whats good may of course align.

This is where it gets tricky to discuss in abstract. In the abstract its reasonable to say “don’t add canonicalization that only aids one pass/set of assumptions”. But that one pass may be indicative of class, its just the only one of that class upstream/visible.

I would want to see which patterns are being referred to here to aid the discussion, esp as we are discussing examples here/in creation of doc/policy. I recall there were some load bearing dim 1 things, and those seem to point to losing information / attaching additional semantics not part of the op world. The vector dialect does need a lot of TLC unfortunately. But it has only 22 canonicalizers so that’s tractable to list & eval. And perhaps that’s a way to make this concrete: use this as example for fine tuning guide lines here.

16

RE ‘do no harm’ and other rules of thumb, I made a brain dump a few months ago around when we discussed vector.shape_cast with @qed and @mehdi_amini and shared it here: How we canonicalize - Google Docs.

The main idea was to assume an abstract cost model that ultimately cares about the number of bitflips and values in the IR. A good candidate for a canon pattern would be something that minimizes either of these across any reasonable target. And conversely, something that increases the number of bitflips on a specific target would most likely make a bad canon pattern. This seems easier with lower-level dialects, which matches my general experience with canon patterns across mlir / llvm / spir-v.

17

I think I have created an unintentional tangent… If so, I retract my comments and will abstain from the discussion from now on.

Here’s my thinking:

and

Indeed. We are talking about different things, and I think this is part of the problem, specifically to MLIR.

To me, single operation canonical form is a trivial thing, unrelated to correctness, and does not impact complex matchers I am worried about above.

MLIR has dialects, which have canonical forms for their ops (your second point), but not for the interaction between their ops (first point). If we want to cover “sub-graph canonical form”, we can only change the dialect documentation and we’re good. For example, in Linalg, a binop(bcast(a), b) is equivalent to a generic(a, b) with bcast map { binop } and one of them is a canonical form. Then it’s easy to have a canonicalizer pass that converts from the non-canonical to the canonical.

MLIR also has interfaces, which multiple dialects can implement. We cannot talk about canonical forms on interface level, but if the ops implementing them have different canonical forms, then passes that operate on interfaces will not match the patterns on both dialects, even if they do implement the same interfaces. This is a more generic version of sub-graph canonical form, from an abstraction level.

Discussing canonical forms for individual ops in a dialect is trivial and just needs some local consensus. The consequences are largely localized and don’t impact passes/transforms all that much. Honestly, I don’t spend much time caring about those.

But having gone through trying to implement pattern matchers for Linalg while somehow still trying to make it work to other dialects that also implement the same interfaces (ex. TilingInterface), I really wanted to have at least a way to describe what is a sub-graph canonical form from a particular point of view that might be different depending on which dialect implements the same interface.

If this is not the point here, then just ignore my previous comments.

18

I think this biases to pure codegen at cost of analysis or transformation. And so makes it less useful for higher level applications. It also biases to CISC ops. While good to avoid introducing a ton of ops, doing fully would limit CSE and general optimization opportunities (e.g., the op that has a fused epilogue attribute that indicates which ops have been fused in as string is now preferred form, so one could argue for local ops fusion as part of canonical form making those epilogues opaque to anything not bespoke). Which is fine for a leaf dialect or project, but limits reuse.

19

You could also apply it to decide whether canonicalization for higher-level ops like vector.shape_cast could results in (or increase) data movement. Also, I think it’s easier to argue that patterns that preserve the abstraction level make sense for canonicalization vs. patterns that, arguably, perform lowering/decomposition.

As a rule of thumb, this ‘abstract cost model’ test allows you to tell if a canon pattern is ‘bad’, but it doesn’t necessarily help you to decide if a pattern is ‘good’.

20

This contradicts one of the current Canonicalizer guidelines:
Pass pipelines should not rely on the canonicalizer pass for correctness. They should work correctly with all instances of the canonicalization pass removed.

In your example, the canonicalizer pass is not optional. Without it, the compilation pipeline cannot make progress. Personally, I also see “reduction of number of cases to pattern match” as one of the main benefits. But the current canonicalizer pass cannot be safely used to do that. (Because it may abort the greedy pattern rewrite without producing a pass failure.)