[mlir][Transforms] Dialect Conversion: Fix folder implementation (#150775)

Operation folders can do two things:

1. Modify IR (in-place op modification). Failing to legalize an in-place
folded operation does not trigger an immediate rollback. This happens
only if the driver decides to try a different lowering path, requiring
it to roll back a bunch of modifications, including the application of
the folder.
2. Create new IR (constant op materialization of a folded attribute).
Failing to legalize a newly created constant op triggers an immediate
rollback.

In-place op modifications should be guarded by
`startOpModification`/`finalizeOpModification` because they are no
different from other in-place op modifications. (They just happen
outside of a pattern, but that does not mean that we should not track
those changes; we are tracking everything else.) This commit adds those
two function calls.

This commit also moves the `rewriter.replaceOp(op, replacementValues);`
function call before the loop nest that legalizes the newly created
constant ops (and therefore `replacementValues`). Conceptually, the
folded op must be replaced before attempting to legalize the constants
because the constant ops may themselves be replaced as part of their own
legalization process. The previous implementation happened to work in
the current conversion driver, but is incompatible with the One-Shot
Dialect Conversion driver, which expects to see the most recent IR at
all time.

From an end-user perspective, this commit should be NFC. A common
folder-rollback pattern that is exercised by multiple tests cases: A
`memref.dim` is folded to `arith.constant`, but `arith.constant` is not
marked as legal as per the conversion target, triggering a rollback.

Note: Folding is generally unsafe in a dialect conversion (see #92683),
but that's a different issue. (In a One-Shot Dialect Conversion, it will
no longer be unsafe.)

Author: matthias-springer

mlir/lib/Transforms/Utils/DialectConversion.cpp

@@ -17,6 +17,7 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Rewrite/PatternApplicator.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -2240,23 +2241,39 @@ OperationLegalizer::legalizeWithFold(Operation *op,
     rewriterImpl.logger.startLine() << "* Fold {\n";
     rewriterImpl.logger.indent();
   });
-  (void)rewriterImpl;
+
+  // Clear pattern state, so that the next pattern application starts with a
+  // clean slate. (The op/block sets are populated by listener notifications.)
+  auto cleanup = llvm::make_scope_exit([&]() {
+    rewriterImpl.patternNewOps.clear();
+    rewriterImpl.patternModifiedOps.clear();
+    rewriterImpl.patternInsertedBlocks.clear();
+  });
+
+  // Upon failure, undo all changes made by the folder.
+  RewriterState curState = rewriterImpl.getCurrentState();
 
   // Try to fold the operation.
   StringRef opName = op->getName().getStringRef();
   SmallVector<Value, 2> replacementValues;
   SmallVector<Operation *, 2> newOps;
   rewriter.setInsertionPoint(op);
+  rewriter.startOpModification(op);
   if (failed(rewriter.tryFold(op, replacementValues, &newOps))) {
     LLVM_DEBUG(logFailure(rewriterImpl.logger, "unable to fold"));
+    rewriter.cancelOpModification(op);
     return failure();
   }
+  rewriter.finalizeOpModification(op);
 
   // An empty list of replacement values indicates that the fold was in-place.
   // As the operation changed, a new legalization needs to be attempted.
   if (replacementValues.empty())
     return legalize(op, rewriter);
 
+  // Insert a replacement for 'op' with the folded replacement values.
+  rewriter.replaceOp(op, replacementValues);
+
   // Recursively legalize any new constant operations.
   for (Operation *newOp : newOps) {
     if (failed(legalize(newOp, rewriter))) {
@@ -2269,16 +2286,12 @@ OperationLegalizer::legalizeWithFold(Operation *op,
             "op '" + opName +
             "' folder rollback of IR modifications requested");
       }
-      // Legalization failed: erase all materialized constants.
-      for (Operation *op : newOps)
-        rewriter.eraseOp(op);
+      rewriterImpl.resetState(
+          curState, std::string(op->getName().getStringRef()) + " folder");
       return failure();
     }
   }
 
-  // Insert a replacement for 'op' with the folded replacement values.
-  rewriter.replaceOp(op, replacementValues);
-
   LLVM_DEBUG(logSuccess(rewriterImpl.logger, ""));
   return success();
 }

mlir/test/Transforms/test-legalize-type-conversion.mlir

@@ -104,8 +104,8 @@ func.func @test_signature_conversion_no_converter() {
   "test.signature_conversion_no_converter"() ({
   // expected-error@below {{failed to legalize unresolved materialization from ('f64') to ('f32') that remained live after conversion}}
   ^bb0(%arg0: f32):
-    "test.type_consumer"(%arg0) : (f32) -> ()
     // expected-note@below{{see existing live user here}}
+    "test.type_consumer"(%arg0) : (f32) -> ()
     "test.return"(%arg0) : (f32) -> ()
   }) : () -> ()
   return