[mlir] GreedyPatternRewriteDriver: Fix termination criteria in OpPatternRewriteDriver

[springerm]

Summary

This driver should iterate until convergence or until the specified op was erased. However, it used to stop when any op was erased.

Depends On: D141920

[springerm]

Note: I found that some downstream project tests (in XLA GPU) fail with this fix. Probably due to broken rewrite patterns.

[springerm]

In D141921#4058600, @springerm wrote:

Note: I found that some downstream project tests (in XLA GPU) fail with this fix. Probably due to broken rewrite patterns.

These were indeed due to broken rewrite patterns that in some cases modified IR without going through the rewriter.

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

494

Can you move this just before opErasedViaPatternRewrites? These are kind of working together.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

public:
   LogicalResult simplifyLocally(Operation *op, int64_t maxNumRewrites,
                                 bool &erased);
 
   // These are hooks implemented for PatternRewriter.
 protected:
   /// If an operation is about to be removed, mark it so that we can let clients
   /// know.
   void notifyOperationRemoved(Operation *op) override {
+    if (this->op == op)
       opErasedViaPatternRewrites = true;
   }
 
   // When a root is going to be replaced, its removal will be notified as well.
   // So there is nothing to do here.
   void notifyRootReplaced(Operation *op, ValueRange replacement) override {}
 
 private:
+  /// Op that is being processed.
+  Operation *op = nullptr;
+
   /// The low-level pattern applicator.
   PatternApplicator matcher;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
   /// Set to true if the operation has been erased via pattern rewrites.
   bool opErasedViaPatternRewrites = false;
 };
 
 } // namespace
 
 /// Performs the rewrites and folding only on `op`. The simplification
 /// converges if the op is erased as a result of being folded, replaced, or
 /// becoming dead, or no more changes happen in an iteration. Returns success if
 /// the rewrite converges in `maxNumRewrites`. `erased` is set to true if `op`
 /// gets erased.
 LogicalResult OpPatternRewriteDriver::simplifyLocally(Operation *op,
                                                       int64_t maxNumRewrites,
                                                       bool &erased) {
+  this->op = op;
   bool changed = false;
   erased = false;
   opErasedViaPatternRewrites = false;
   int64_t numRewrites = 0;
   // Iterate until convergence or until maxNumRewrites. Deletion of the op as
   // a result of being dead or folded is convergence.
   do {
     if (numRewrites >= maxNumRewrites &&