legalizeWithFold vs legalizeWithCanonicalization

MLIR
#1

Hello!

In my dialect, I have an Op FooOp that has a folder and a canonicalization pattern (i.e., hasFolder = 1 and hasCanonicalizer = 1). Both of them work together (in the canonicalization pass) to produce the canonical form of FooOp. The folder alone will leave FooOp in a partially canonicalized state that is awkward for my subsequent dialect conversion passes.

A dialect conversion pass will attempt to “legalize” Ops using the OperationLegalizer. The legalizer will call a method called legalizeWithFold to try folding the Ops. As you can guess, that leaves FooOp in an awkward partially canonicalized state.

Of course, I can always run the canonicalizer pass before dialect conversion and that’s fine. However, I got curious:

Would it make sense to have a legalizeWithCanonicalization method instead ?
Such method will basically run the full canonicalization pass before dialect conversions.

Thank you in advance !

#2

Sorry for the late response.

Legalizing with canonicalization patterns is something that I have considered in the past, but there have been several blockers that have prevented that from happening:

  • Canonicalization patterns have traditionally done things that DialectConversion can’t support
    For a really long time, a lot of canonicalization patterns didn’t use the PatternRewriter for performing updates to the IR. As you might case, this effectively broke DialectConversion given that it heavily requires on the PatternRewriter API being properly used. This is less of an issue now that I started more heavily enforcing that the PatternRewriter be used always within a pattern, but there are likely some lingering issues in this area.
  • Dialect conversion has a large amount of intermixed IR at various levels, which doesn’t always play nicely with canonicalizations
    Legalizing with the fold method has always been fine because the API contract makes it local to the current operation. Canonicalizations may make assumptions about how the IR is structured, which work out fine normally, but break given the intermixed nature of the IR during dialect conversion. e.g. The parent of one operation may not be what it traditionally expects, e.g. a SPV function may have already been converted to an LLVM function by the time that the inner operations are converted.
  • The pattern list is already frozen when DialectConversion runs
    With PDL on the rise, caching/sharing patterns is extremely important given that the cost of recompiling a set of PDL patterns is not cheap. If DialectConversion grabbed the canonicalization patterns automatically, it would need to recompile them each run of the pass. Not even considering PDL, just grabbing canonicalization patterns from the context has shown up on profiles. The other alternative would be that the user provides them to dialect conversion. Though at this point, the user could just preemptively add them to their pattern list and dialect conversion wouldn’t have to know about them specifically at all.

– River