Continuing our conversation from this weeks ODM, we had an internal meeting today and came up with an algorithm and implementation thoughts which wouldn’t require modifications to the CIRCT optimizations which use the PatternWriters. I’d like to get confirmation that the algorithm is sound and that it could be implemented the way I think it can be.
Problem statement
Given code a, its optimized version A, and a patched version of a called a', generate optimized code for a' (which we’ll call A') with a minimal number of changes between A and A', both in terms of netlist and SystemVerilog text. We’re also given a mapping between operations in a and a' for all operations which are present in both (have not changed). We also assume that we can call a given transformation on an operation for only that operation (as we can for the canonicalizers) We further assume a deterministic compiler.
Algorithm
Run a given transformation on a (to produce A). For each transform on each operation, log the operation and the transformation in the order in which they are run.
Iterate through the operations in a' in the same order as the transform log for a (for all ops which exist in both) and run the optimization (not the transformation) on each operation, monitoring the transformation which it does on a'. If and only if the same optimization was applied to a, allow it to persist. If the transformation was not applied, roll back the transformation.
Profit.
The idea is to apply as many of the transforms to a' which were applied to a which are still valid to apply. By running the optimizations on a' rather than just blindly applying them, we ensure that they are still valid. By checking that a transformation was applied to a, we ensure that differences between a and a' do not cause non-linear changes to occur.
As a memory optimization step, we can run the steps concurrently.
Implementation
This algorithm requires two capabilities: transformation observation and rollback.
We can implement the observation component by implementing a Listener and recording the transformations.
I seem to recall (but cannot find at the moment) that there is some transformation rewind / rollback functionality which can be applied if transformations occur through a PatternRewriter.
I think this would work on the canonicalizers as is. Further, I’m pretty sure this could be used on any optimization which uses the dialect conversion framework with the addition of a “use this listener/rollback widget” argument.
Looking into Listener, I wouldn’t say that this will out of the box work, and would most likely require upstream changes.
Existing pattern rewriters (e.g. the ConversionPatternRewriter seem to take ownership over the “listening” part of the conversion infrastructure.
OpBuilder::Listener doesn’t seem to have a notion of currently-executing pass.
Ideally, we’d hope those listeners to just be API sinks that can be provided to any pattern rewriter, but that does not seem to have been the design goal/intention, as far as I can tell. Based on the abundance of final specifiers in the pattern rewriters, I neither think they were designed to be subclassed (in case we wanted a CIRCT-specific pattern rewriter - which I do not think we do).
I haven’t looked into this in detail to see what we can do with the current Listener, so I’ll defer to Morten’s analysis. Maybe there is something we can do here. There was an RFC about making Listener more composable: [RFC] Composition-based Listener for RewriterBase.
A couple other ideas that seem related:
There was a proposal to add a new fine-grained IR listener: [RFC] Introduce the concept of IR listeners in MLIR. That might be a really good fit for this if upstream MLIR goes this route. Someone even commented about “diffing” IR, which sounds related to this notion of ECOs.
I agree that’s what we’d ideally want. Given that we’re only interested in the canonicalizers, though, I think it’s reasonable to write our own canonicalization pass ‘eco-able-canonicalize’. The greedy rewriter (which the standard canonicalize pass uses) supports adding a listener. And createCanonicalizerPass accepts it.
That ConversionPatternRewriter is final, however, could be a problem since we want rollback capabilities. Ideally, rollback capabilities would be accessible through a public API which a listener (or some other “listener” interface) could call into. So I agree, the best way forward is changing upstream to support our needs.
Yeah, I was initially thinking about using actions. If I understand correctly, I think it’d work for the listening capabilities we’d need. Not clear it it’d be able to expose the rollback call which we’d need.
I’ve been looking into the internals of rollback/commit llvm-project/DialectConversion.cpp at c0967995d254fd2dd24ba74afa46df5f559f11ba · llvm/llvm-project · GitHub. There’s a lot of useful things in there not just for rollback/commit, but (and this makes total sense in retrospect) for tracking the specific changes. This would come in handy for answering the question, “does the transformation on a match the one executed on a'?” as well as “if yes, commit the changes”.