[mlir][Transforms] Add support for ConversionPatternRewriter::replaceAllUsesWith (#155244)

This commit generalizes `replaceUsesOfBlockArgument` to
`replaceAllUsesWith`. In rollback mode, the same restrictions keep
applying: a value cannot be replaced multiple times and a call to
`replaceAllUsesWith` will replace all current and future uses of the
`from` value.

`replaceAllUsesWith` is now fully supported and its behavior is
consistent with the remaining dialect conversion API. Before this
commit, `replaceAllUsesWith` was immediately reflected in the IR when
running in rollback mode. After this commit, `replaceAllUsesWith`
changes are materialized in a delayed fashion, at the end of the dialect
conversion. This is consistent with the `replaceUsesOfBlockArgument` and
`replaceOp` APIs.

`replaceAllUsesExcept` etc. are still not supported and will be
deactivated on the `ConversionPatternRewriter` (when running in rollback
mode) in a follow-up commit.

Note for LLVM integration: Replace `replaceUsesOfBlockArgument` with
`replaceAllUsesWith`. If you are seeing failures, you may have patterns
that use `replaceAllUsesWith` incorrectly (e.g., being called multiple
times on the same value) or bypass the rewriter API entirely. E.g., such
failures were mitigated in Flang by switching to the walk-patterns
driver (#156171).

You can temporarily reactivate the old behavior by calling
`RewriterBase::replaceAllUsesWith`. However, note that that behavior is
faulty in a dialect conversion. E.g., the base
`RewriterBase::replaceAllUsesWith` implementation does not see uses of
the `from` value that have not materialized yet and will, therefore, not
replace them.

Author: matthias-springer

flang/lib/Optimizer/OpenMP/DoConcurrentConversion.cpp

@@ -444,6 +444,18 @@ class DoConcurrentConversion
   mlir::SymbolTable &moduleSymbolTable;
 };
 
+/// A listener that forwards notifyOperationErased to the given callback.
+struct CallbackListener : public mlir::RewriterBase::Listener {
+  CallbackListener(std::function<void(mlir::Operation *op)> onOperationErased)
+      : onOperationErased(onOperationErased) {}
+
+  void notifyOperationErased(mlir::Operation *op) override {
+    onOperationErased(op);
+  }
+
+  std::function<void(mlir::Operation *op)> onOperationErased;
+};
+
 class DoConcurrentConversionPass
     : public flangomp::impl::DoConcurrentConversionPassBase<
           DoConcurrentConversionPass> {
@@ -468,6 +480,10 @@ class DoConcurrentConversionPass
     }
 
     llvm::DenseSet<fir::DoConcurrentOp> concurrentLoopsToSkip;
+    CallbackListener callbackListener([&](mlir::Operation *op) {
+      if (auto loop = mlir::dyn_cast<fir::DoConcurrentOp>(op))
+        concurrentLoopsToSkip.erase(loop);
+    });
     mlir::RewritePatternSet patterns(context);
     patterns.insert<DoConcurrentConversion>(
         context, mapTo == flangomp::DoConcurrentMappingKind::DCMK_Device,
@@ -480,8 +496,11 @@ class DoConcurrentConversionPass
     target.markUnknownOpDynamicallyLegal(
         [](mlir::Operation *) { return true; });
 
-    if (mlir::failed(
-            mlir::applyFullConversion(module, target, std::move(patterns)))) {
+    mlir::ConversionConfig config;
+    config.allowPatternRollback = false;
+    config.listener = &callbackListener;
+    if (mlir::failed(mlir::applyFullConversion(module, target,
+                                               std::move(patterns), config))) {
       signalPassFailure();
     }
   }

mlir/include/mlir/IR/PatternMatch.h

@@ -633,7 +633,7 @@ class RewriterBase : public OpBuilder {
 
   /// Find uses of `from` and replace them with `to`. Also notify the listener
   /// about every in-place op modification (for every use that was replaced).
-  void replaceAllUsesWith(Value from, Value to) {
+  virtual void replaceAllUsesWith(Value from, Value to) {
     for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
       Operation *op = operand.getOwner();
       modifyOpInPlace(op, [&]() { operand.set(to); });

mlir/include/mlir/Transforms/DialectConversion.h

@@ -854,15 +854,29 @@ class ConversionPatternRewriter final : public PatternRewriter {
       Region *region, const TypeConverter &converter,
       TypeConverter::SignatureConversion *entryConversion = nullptr);
 
-  /// Replace all the uses of the block argument `from` with `to`. This
-  /// function supports both 1:1 and 1:N replacements.
+  /// Replace all the uses of `from` with `to`. The type of `from` and `to` is
+  /// allowed to differ. The conversion driver will try to reconcile all type
+  /// mismatches that still exist at the end of the conversion with
+  /// materializations. This function supports both 1:1 and 1:N replacements.
   ///
-  /// Note: If `allowPatternRollback` is set to "true", this function replaces
-  /// all current and future uses of the block argument. This same block
-  /// block argument must not be replaced multiple times. Uses are not replaced
-  /// immediately but in a delayed fashion. Patterns may still see the original
-  /// uses when inspecting IR.
-  void replaceUsesOfBlockArgument(BlockArgument from, ValueRange to);
+  /// Note: If `allowPatternRollback` is set to "true", this function behaves
+  /// slightly different:
+  ///
+  /// 1. All current and future uses of `from` are replaced. The same value must
+  ///    not be replaced multiple times. That's an API violation.
+  /// 2. Uses are not replaced immediately but in a delayed fashion. Patterns
+  ///    may still see the original uses when inspecting IR.
+  /// 3. Uses within the same block that appear before the defining operation
+  ///    of the replacement value are not replaced. This allows users to
+  ///    perform certain replaceAllUsesExcept-style replacements, even though
+  ///    such API is not directly supported.
+  ///
+  /// Note: In an attempt to align the ConversionPatternRewriter and
+  /// RewriterBase APIs, (3) may be removed in the future.
+  void replaceAllUsesWith(Value from, ValueRange to);
+  void replaceAllUsesWith(Value from, Value to) override {
+    replaceAllUsesWith(from, ValueRange{to});
+  }
 
   /// Return the converted value of 'key' with a type defined by the type
   /// converter of the currently executing pattern. Return nullptr in the case

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

@@ -284,7 +284,7 @@ static void restoreByValRefArgumentType(
         cast<TypeAttr>(byValRefAttr->getValue()).getValue());
 
     Value valueArg = LLVM::LoadOp::create(rewriter, arg.getLoc(), resTy, arg);
-    rewriter.replaceUsesOfBlockArgument(arg, valueArg);
+    rewriter.replaceAllUsesWith(arg, valueArg);
   }
 }