Exposing transform mutation tracking and "rollback" capabilities

MLIR
1

Original, application specific discussion: CIRCT ECO in the presence of transformations. Not a necessary read for understanding this post.

We have an interesting use case for patten rewriters: we need to monitor IR mutations on one program a, then do the same thing on a slightly modified version of a (a') and only apply the IR mutations to a' which were applied to a. Let’s assume we have already derived a mapping of Operations in a to Operations in a'. We define “mutation” as a single call to a canonicalizer or a matchAndRewrite method.

As I see it, this requires the following capabilities:

  1. To “listen” to the mutations and record what individual calls into the rewriter are doing.
  2. To compare the records of two mutations to determine if they are the same.
  3. To selectively apply or discard the mutation on a' based on said comparison.

Digging around the MLIR code, I see that #1 and #3 already exist in DialectConversion.cpp. Specifically in ConversionPatternRewriterImpl. Adding the ability to check equality of two of these given an Operation mapping should be relatively straightforward.

Since our requirements are somewhat domain-specific, I don’t expect it to be supported by MLIR proper. IMO, it should live in the CIRCT project (on which I work). It would be ideal to reuse all the bookkeeping in MLIR’s dialect conversion framework which is currently private. What do people think of making the functionality of ConversionPatternRewriterImpl a public API and being able to subclass it (for adding a comparison method, for instance)?

I see the functionality as generally useful, which implies it should be publicly accessible. I also see, however, that code in DialectConversion.cpp tends to interact with ConversionPatternRewriterImpl internal data structures, so just moving it to a header file probably wouldn’t be acceptable.

2

Have you seen [RFC] Introduce the concept of IR listeners in MLIR ?

3

I had read it in the past, but I just re-skimmed the thread.

It seems to me that this is a mechanism for listening and – for my purposes – not materially different to a PatternRewriter listener. It doesn’t get us change tracking/accounting (we’d have to essentially replicate what ConversionPatternRewriterImpl does). It also doesn’t get us (in fact it impedes) the apply/discardRewrites functionality there. (It would impede it since we don’t get the option to delay rewrites until we’ve decided to commit the changes.) All of the passes we’re immediately concerned with use a PattenRewriter so the generality is not immediately necessary.