[mlir] [mem2reg] Adapt to be pattern-friendly.

[Moxinilian]

Summary

This revision modifies the mem2reg interfaces and algorithm to be more
omfortable to use as a pattern. The motivation behind this is that
currently the pattern needs to be applied to the scope op of the region
in which allocators should be promoted. However, a more natural way to
apply the pattern would be to apply it on the allocator directly. This
is not only clearer but easier to parallelize.

This revision changes the mem2reg pattern to operate this way. This
required restraining the interfaces to only mutate IR using
RewriterBase, as the previously used escape hatch is not granular enough
to match on the region that is modified only. This has the unfortunate
cost of preventing batching allocator promotion and making the block
argument adding logic more complex. Because batching no longer made any
sense, I made the internal analyzer/promoter decoupling private again.

This also adds statistics to the mem2reg infrastructure.

[gysit]

Nice! I have some minor comments.

mlir/lib/Transforms/Mem2Reg.cpp

455–482

nit: back() may be more elegant here?

476–478

I believe drop_back may allow us to save the copy here?

630

Can we enable the region simplification and add a pass option to disable it during testing?

The canonicalize pass already has such a pass flag:

Option<"enableRegionSimplification", "region-simplify", "bool",
       /*default=*/"true",
       "Perform control flow optimizations to the region tree">,

633

nit: You can pass the op directly to the rewriter and the rewriter will iterate all regions, i.e.:

(void)applyPatternsAndFoldGreedily(scopeOp, frozen, config)

[Moxinilian]

Fix forgotten direct IR mutation + fail if patterns do not converge.

[Moxinilian]

Address comments + add statistics to mem2reg.

mlir/lib/Transforms/Mem2Reg.cpp

633

Am I allowed to do this? What if the scopeOp is an allocator and I delete it? Wouldn't that break the isolation constraint of passes?

[gysit]

mlir/lib/Transforms/Mem2Reg.cpp

633

The patterns apply only to the regions and not to the scope op, see implementation:

inline LogicalResult applyPatternsAndFoldGreedily(
    Operation *op, const FrozenRewritePatternSet &patterns,
    GreedyRewriteConfig config = GreedyRewriteConfig()) {
  bool failed = false;
  for (Region &region : op->getRegions())
    failed |= applyPatternsAndFoldGreedily(region, patterns, config).failed();
  return failure(failed);
}

[gysit]

Nice cleanup!

Can you add a small test case for the statistics so that I can see it in action?

mlir/include/mlir/Transforms/Mem2Reg.h

21

nit: I would probably use mlir::Pass::Statistic here?

mlir/lib/Transforms/Mem2Reg.cpp

623

If I understand correctly Mem2RegPattern stores a reference to Mem2RegStatistics? Are you sure this is not resulting in a use-after-free. A temporary is normally only valid during the function execution - here the construction of Mem2RegPattern - but not later on once the pattern is applied?

Allocating Mem2RegPattern on the stack of runOnOperation seems the safe choice here?

[Moxinilian]

Address comments.

mlir/include/mlir/Transforms/Mem2Reg.h

21

I feel like "Pass::Statistic" has a connotation for passes. What if you want to have your own custom statistics? Pass::Statistic inherits from llvm::Statistic anyway so there is no practical advantage.

[gysit]

LGTM!

[Moxinilian]

Fix memory issue with default constructor.

mlir/include/mlir/Interfaces/MemorySlotInterfaces.h

 //===-- Mem2RegInterfaces.h - Mem2Reg interfaces ----------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_MEMORYSLOTINTERFACES_H
 #define MLIR_INTERFACES_MEMORYSLOTINTERFACES_H
 
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 
 /// Represents a slot in memory. This is generated by an allocating operation
 /// (for example alloca).
 struct MemorySlot {
   /// Pointer to the memory slot, used by operations to refer to it.
   Value ptr;

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

InterfaceMethod<[{
         whether the slot pointer is properly used or not. This allocator is the
         "owner" of the returned slots, meaning no two allocators should return
         the same slot. The content of the memory slot must only be reachable
         using loads and stores to the provided slot pointer, no aliasing is
         allowed.
 
         Promotion of the slot will lead to the slot pointer no longer being
         used, leaving the content of the memory slot unreachable.
+
+        No IR mutation is allowed in this method.
       }], "::llvm::SmallVector<::mlir::MemorySlot>", "getPromotableSlots",
       (ins)
     >,
     InterfaceMethod<[{
         Provides the default Value of this memory slot. The provided Value
         will be used as the reaching definition of loads done before any store.
         This Value must outlive the promotion and dominate all the uses of this
-        slot's pointer. The provided builder can be used to create the default
+        slot's pointer. The provided rewriter can be used to create the default
         value on the fly.
 
-        The builder is located at the beginning of the block where the slot
-        pointer is defined.
+        The rewriter is located at the beginning of the block where the slot
+        pointer is defined. All IR mutations must happen through the rewriter.
       }], "::mlir::Value", "getDefaultValue",
-      (ins "const ::mlir::MemorySlot &":$slot, "::mlir::OpBuilder &":$builder)
+      (ins
+        "const ::mlir::MemorySlot &":$slot,
+        "::mlir::RewriterBase &":$rewriter)
     >,
     InterfaceMethod<[{
         Hook triggered for every new block argument added to a block.
         This will only be called for slots declared by this operation.
 
-        The builder is located at the beginning of the block on call.
+        The rewriter is located at the beginning of the block on call. All IR
+        mutations must happen through the rewriter.
       }],
       "void", "handleBlockArgument",
       (ins
         "const ::mlir::MemorySlot &":$slot,
         "::mlir::BlockArgument":$argument,
-        "::mlir::OpBuilder &":$builder
+        "::mlir::RewriterBase &":$rewriter
       )
     >,
     InterfaceMethod<[{
         Hook triggered once the promotion of a slot is complete. This can
         also clean up the created default value if necessary.
         This will only be called for slots declared by this operation.
+
+        All IR mutations must happen through the rewriter.
       }],
       "void", "handlePromotionComplete",
-      (ins "const ::mlir::MemorySlot &":$slot, "::mlir::Value":$defaultValue)
+      (ins
+        "const ::mlir::MemorySlot &":$slot, 
+        "::mlir::Value":$defaultValue,
+        "::mlir::RewriterBase &":$rewriter)
     >,
   ];
 }
 
 def PromotableMemOpInterface : OpInterface<"PromotableMemOpInterface"> {
   let description = [{
     Describes an operation that can load from memory slots and/or store
     to memory slots. Loads and stores must be of whole values of the same
     type as the slot itself.
 
     For a memory operation on a slot to be valid, it must operate on the slot
     pointer *only as a pointer to an element of the type of the slot*.
 
     If the same operation does both loads and stores on the same slot, the
     load must semantically happen first.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
         Gets whether this operation loads from the specified slot.
+
+        No IR mutation is allowed in this method.
       }],
       "bool", "loadsFrom",
       (ins "const ::mlir::MemorySlot &":$slot)
     >,
     InterfaceMethod<[{
         Gets the value stored to the provided memory slot, or returns a null
         value if this operation does not store to this slot. An operation
         storing a value to a slot must always be able to provide the value it
         stores. This method is only called on operations that use the slot.
+
+        No IR mutation is allowed in this method.
       }],
       "::mlir::Value", "getStored",
       (ins "const ::mlir::MemorySlot &":$slot)
     >,
     InterfaceMethod<[{
         Checks that this operation can be promoted to no longer use the provided
         blocking uses, in the context of promoting `slot`.
 
         If the removal procedure of the use will require that other uses get
         removed, that dependency should be added to the `newBlockingUses`
         argument. Dependent uses must only be uses of results of this operation.
+
+        No IR mutation is allowed in this method.
       }], "bool", "canUsesBeRemoved",
       (ins "const ::mlir::MemorySlot &":$slot,
            "const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &":$blockingUses,
            "::llvm::SmallVectorImpl<::mlir::OpOperand *> &":$newBlockingUses)
     >,
     InterfaceMethod<[{
         Transforms IR to ensure that the current operation does not use the
         provided memory slot anymore. `reachingDefinition` contains the value
InterfaceMethod<[{
         After calling this method, the blocking uses should have disappeared
         or this operation should have scheduled its own deletion.
 
         This method will only be called after ensuring promotion is allowed via
         `canUseBeRemoved`. The requested blocking use removal may or may not
         have been done at the point of calling this method, but it will be done
         eventually.
 
-        The builder is located after the promotable operation on call.
+        The rewriter is located after the promotable operation on call. All IR
+        mutations must happen through the rewriter.
       }],
       "::mlir::DeletionKind",
       "removeBlockingUses",
       (ins "const ::mlir::MemorySlot &":$slot,
            "const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
-           "::mlir::OpBuilder &":$builder,
+           "::mlir::RewriterBase &":$rewriter,
            "::mlir::Value":$reachingDefinition)
     >,
   ];
 }
 
 def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
   let description = [{
     Describes an operation that can be transformed or deleted so it no longer
     uses a provided value (blocking use), in case this would allow the promotion
     of a memory slot.
   }];
   let cppNamespace = "::mlir";
 
   let methods = [
     InterfaceMethod<[{
         Checks that this operation can be promoted to no longer use the provided
         blocking uses, in order to allow optimization.
 
         If the removal procedure of the use will require that other uses get
         removed, that dependency should be added to the `newBlockingUses`
         argument. Dependent uses must only be uses of results of this operation.
+
+        No IR mutation is allowed in this method.
       }], "bool", "canUsesBeRemoved",
       (ins "const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &":$blockingUses,
            "::llvm::SmallVectorImpl<::mlir::OpOperand *> &":$newBlockingUses)
     >,
     InterfaceMethod<[{
         Transforms IR to ensure that the current operation does not use the
         provided blocking uses anymore. In contrast to
         `PromotableMemOpInterface`, operations implementing this interface
InterfaceMethod<[{
         After calling this method, the blocking uses should have disappeared
         or this operation should have scheduled its own deletion.
 
         This method will only be called after ensuring promotion is allowed via
         `canUseBeRemoved`. The requested blocking use removal may or may not
         have been done at the point of calling this method, but it will be done
         eventually.
 
-        The builder is located after the promotable operation on call.
+        The rewriter is located after the promotable operation on call. All IR
+        mutations must happen through the rewriter.
       }],
       "::mlir::DeletionKind",
       "removeBlockingUses",
       (ins "const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
-           "::mlir::OpBuilder &":$builder)
+           "::mlir::RewriterBase &":$rewriter)
     >,
   ];
 }
 
 #endif // MLIR_INTERFACES_MEMORYSLOTINTERFACES

mlir/include/mlir/Transforms/Mem2Reg.h

 //===-- Mem2Reg.h - Mem2Reg definitions -------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_MEM2REG_H
 #define MLIR_TRANSFORMS_MEM2REG_H
 
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
+#include "llvm/ADT/Statistic.h"
 
 namespace mlir {
 
-/// Information computed during promotion analysis used to perform actual
-/// promotion.
-struct MemorySlotPromotionInfo {
-  /// Blocks for which at least two definitions of the slot values clash.
-  SmallPtrSet<Block *, 8> mergePoints;
-  /// Contains, for each operation, which uses must be eliminated by promotion.
-  /// This is a DAG structure because if an operation must eliminate some of
-  /// its uses, it is because the defining ops of the blocking uses requested
-  /// it. The defining ops therefore must also have blocking uses or be the
-  /// starting point of the bloccking uses.
-  DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> userToBlockingUses;
-};
-
-/// Computes information for basic slot promotion. This will check that direct
-/// slot promotion can be performed, and provide the information to execute the
-/// promotion. This does not mutate IR.
-class MemorySlotPromotionAnalyzer {
-public:
-  MemorySlotPromotionAnalyzer(MemorySlot slot, DominanceInfo &dominance)
-      : slot(slot), dominance(dominance) {}
-
-  /// Computes the information for slot promotion if promotion is possible,
-  /// returns nothing otherwise.
-  std::optional<MemorySlotPromotionInfo> computeInfo();
-
-private:
-  /// Computes the transitive uses of the slot that block promotion. This finds
-  /// uses that would block the promotion, checks that the operation has a
-  /// solution to remove the blocking use, and potentially forwards the analysis
-  /// if the operation needs further blocking uses resolved to resolve its own
-  /// uses (typically, removing its users because it will delete itself to
-  /// resolve its own blocking uses). This will fail if one of the transitive
-  /// users cannot remove a requested use, and should prevent promotion.
-  LogicalResult computeBlockingUses(
-      DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> &userToBlockingUses);
-
-  /// Computes in which blocks the value stored in the slot is actually used,
-  /// meaning blocks leading to a load. This method uses `definingBlocks`, the
-  /// set of blocks containing a store to the slot (defining the value of the
-  /// slot).
-  SmallPtrSet<Block *, 16>
-  computeSlotLiveIn(SmallPtrSetImpl<Block *> &definingBlocks);
-
-  /// Computes the points in which multiple re-definitions of the slot's value
-  /// (stores) may conflict.
-  void computeMergePoints(SmallPtrSetImpl<Block *> &mergePoints);
-
-  /// Ensures predecessors of merge points can properly provide their current
-  /// definition of the value stored in the slot to the merge point. This can
-  /// notably be an issue if the terminator used does not have the ability to
-  /// forward values through block operands.
-  bool areMergePointsUsable(SmallPtrSetImpl<Block *> &mergePoints);
-
-  MemorySlot slot;
-  DominanceInfo &dominance;
-};
-
-/// The MemorySlotPromoter handles the state of promoting a memory slot. It
-/// wraps a slot and its associated allocator. This will perform the mutation of
-/// IR.
-class MemorySlotPromoter {
-public:
-  MemorySlotPromoter(MemorySlot slot, PromotableAllocationOpInterface allocator,
-                     OpBuilder &builder, DominanceInfo &dominance,
-                     MemorySlotPromotionInfo info);
-
-  /// Actually promotes the slot by mutating IR. Promoting a slot does not
-  /// invalidate the MemorySlotPromotionInfo of other slots.
-  void promoteSlot();
-
-private:
-  /// Computes the reaching definition for all the operations that require
-  /// promotion. `reachingDef` is the value the slot should contain at the
-  /// beginning of the block. This method returns the reached definition at the
-  /// end of the block.
-  Value computeReachingDefInBlock(Block *block, Value reachingDef);
-
-  /// Computes the reaching definition for all the operations that require
-  /// promotion. `reachingDef` corresponds to the initial value the
-  /// slot will contain before any write, typically a poison value.
-  void computeReachingDefInRegion(Region *region, Value reachingDef);
-
-  /// Removes the blocking uses of the slot, in topological order.
-  void removeBlockingUses();
-
-  /// Lazily-constructed default value representing the content of the slot when
-  /// no store has been executed. This function may mutate IR.
-  Value getLazyDefaultValue();
-
-  MemorySlot slot;
-  PromotableAllocationOpInterface allocator;
-  OpBuilder &builder;
-  /// Potentially non-initialized default value. Use `getLazyDefaultValue` to
-  /// initialize it on demand.
-  Value defaultValue;
-  /// Contains the reaching definition at this operation. Reaching definitions
-  /// are only computed for promotable memory operations with blocking uses.
-  DenseMap<PromotableMemOpInterface, Value> reachingDefs;
-  DominanceInfo &dominance;
-  MemorySlotPromotionInfo info;
+struct Mem2RegStatistics {
+  llvm::Statistic *promotedAmount = nullptr;
+  llvm::Statistic *newBlockArgumentAmount = nullptr;
 };
 
 /// Pattern applying mem2reg to the regions of the operations on which it
 /// matches.
-class Mem2RegPattern : public RewritePattern {
+class Mem2RegPattern
+    : public OpInterfaceRewritePattern<PromotableAllocationOpInterface> {
 public:
-  using RewritePattern::RewritePattern;
+  using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
 
-  Mem2RegPattern(MLIRContext *ctx, PatternBenefit benefit = 1)
-      : RewritePattern(MatchAnyOpTypeTag(), benefit, ctx) {}
+  Mem2RegPattern(MLIRContext *context, const Mem2RegStatistics &statistics = {},
+                 PatternBenefit benefit = 1)
+      : OpInterfaceRewritePattern(context, benefit), statistics(statistics) {}
 
-  LogicalResult matchAndRewrite(Operation *op,
+  LogicalResult matchAndRewrite(PromotableAllocationOpInterface allocator,
                                 PatternRewriter &rewriter) const override;
+
+private:
+  const Mem2RegStatistics &statistics;
 };
 
 /// Attempts to promote the memory slots of the provided allocators. Succeeds if
 /// at least one memory slot was promoted.
 LogicalResult
 tryToPromoteMemorySlots(ArrayRef<PromotableAllocationOpInterface> allocators,
-                        OpBuilder &builder, DominanceInfo &dominance);
+                        RewriterBase &rewriter,
+                        const Mem2RegStatistics &statistics = {});
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_MEM2REG_H

mlir/include/mlir/Transforms/Passes.td

let description = [{
     the memory slot reach operations that use the memory slot pointer. It
     will rewire or remove operations that use the slot pointer so they no
     longer use it. If any of this is not possible, the IR will be left
     without mutation.
 
     This pass only supports unstructured control-flow. Promotion of operations
     within subregions will not happen.
   }];
+
+  let options = [
+    Option<"enableRegionSimplification", "region-simplify", "bool",
+       /*default=*/"true",
+       "Perform control flow optimizations to the region tree">,
+  ];
+
+  let statistics = [
+    Statistic<"promotedAmount",
+              "promoted slots",
+              "Number of promoted memory slot">,
+    Statistic<"newBlockArgumentAmount",
+              "new block args",
+              "Total number of block arguments added">,
+  ];
 }
 
 def PrintOpStats : Pass<"print-op-stats"> {
   let summary = "Print statistics of operations";
   let constructor = "mlir::createPrintOpStatsPass()";
   let options = [
     Option<"printAsJSON", "json", "bool", /*default=*/"false",
            "print the stats as JSON">

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp

 llvm::SmallVector<MemorySlot> LLVM::AllocaOp::getPromotableSlots() {
   if (!getOperation()->getBlock()->isEntryBlock())
     return {};
 
   return {MemorySlot{getResult(), getAllocaElementType(*this)}};
 }
 
 Value LLVM::AllocaOp::getDefaultValue(const MemorySlot &slot,
-                                      OpBuilder &builder) {
-  return builder.create<LLVM::UndefOp>(getLoc(), slot.elemType);
+                                      RewriterBase &rewriter) {
+  return rewriter.create<LLVM::UndefOp>(getLoc(), slot.elemType);
 }
 
 void LLVM::AllocaOp::handleBlockArgument(const MemorySlot &slot,
                                          BlockArgument argument,
-                                         OpBuilder &builder) {
+                                         RewriterBase &rewriter) {
   for (Operation *user : getOperation()->getUsers())
     if (auto declareOp = llvm::dyn_cast<LLVM::DbgDeclareOp>(user))
-      builder.create<LLVM::DbgValueOp>(declareOp.getLoc(), argument,
+      rewriter.create<LLVM::DbgValueOp>(declareOp.getLoc(), argument,
                                         declareOp.getVarInfo());
 }
 
 void LLVM::AllocaOp::handlePromotionComplete(const MemorySlot &slot,
-                                             Value defaultValue) {
+                                             Value defaultValue,
+                                             RewriterBase &rewriter) {
   if (defaultValue && defaultValue.use_empty())
-    defaultValue.getDefiningOp()->erase();
-  erase();
+    rewriter.eraseOp(defaultValue.getDefiningOp());
+  rewriter.eraseOp(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // Interfaces for LoadOp/StoreOp
 //===----------------------------------------------------------------------===//
 
 bool LLVM::LoadOp::loadsFrom(const MemorySlot &slot) {
   return getAddr() == slot.ptr;
bool LLVM::LoadOp::canUsesBeRemoved(
   // be removed (provided it loads the exact stored value and is not
   // volatile).
   return blockingUse == slot.ptr && getAddr() == slot.ptr &&
          getResult().getType() == slot.elemType && !getVolatile_();
 }
 
 DeletionKind LLVM::LoadOp::removeBlockingUses(
     const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder, Value reachingDefinition) {
+    RewriterBase &rewriter, Value reachingDefinition) {
   // `canUsesBeRemoved` checked this blocking use must be the loaded slot
   // pointer.
-  getResult().replaceAllUsesWith(reachingDefinition);
+  rewriter.replaceAllUsesWith(getResult(), reachingDefinition);
   return DeletionKind::Delete;
 }
 
 bool LLVM::StoreOp::canUsesBeRemoved(
     const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   if (blockingUses.size() != 1)
     return false;
   Value blockingUse = (*blockingUses.begin())->get();
   // If the blocking use is the slot ptr itself, dropping the store is
   // fine, provided we are currently promoting its target value. Don't allow a
   // store OF the slot pointer, only INTO the slot pointer.
   return blockingUse == slot.ptr && getAddr() == slot.ptr &&
          getValue() != slot.ptr && getValue().getType() == slot.elemType &&
          !getVolatile_();
 }
 
 DeletionKind LLVM::StoreOp::removeBlockingUses(
     const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder, Value reachingDefinition) {
+    RewriterBase &rewriter, Value reachingDefinition) {
   // `canUsesBeRemoved` checked this blocking use must be the stored slot
   // pointer.
   for (Operation *user : slot.ptr.getUsers())
     if (auto declareOp = llvm::dyn_cast<LLVM::DbgDeclareOp>(user))
-      builder.create<LLVM::DbgValueOp>(declareOp->getLoc(), getValue(),
+      rewriter.create<LLVM::DbgValueOp>(declareOp->getLoc(), getValue(),
                                         declareOp.getVarInfo());
   return DeletionKind::Delete;
 }
 
 //===----------------------------------------------------------------------===//
 // Interfaces for discardable OPs
 //===----------------------------------------------------------------------===//
 
 /// Conditions the deletion of the operation to the removal of all its uses.
 static bool forwardToUsers(Operation *op,
                            SmallVectorImpl<OpOperand *> &newBlockingUses) {
   for (Value result : op->getResults())
     for (OpOperand &use : result.getUses())
       newBlockingUses.push_back(&use);
   return true;
 }
 
 bool LLVM::BitcastOp::canUsesBeRemoved(
     const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return forwardToUsers(*this, newBlockingUses);
 }
 
 DeletionKind LLVM::BitcastOp::removeBlockingUses(
-    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, RewriterBase &rewriter) {
   return DeletionKind::Delete;
 }
 
 bool LLVM::AddrSpaceCastOp::canUsesBeRemoved(
     const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return forwardToUsers(*this, newBlockingUses);
 }
 
 DeletionKind LLVM::AddrSpaceCastOp::removeBlockingUses(
-    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, RewriterBase &rewriter) {
   return DeletionKind::Delete;
 }
 
 bool LLVM::LifetimeStartOp::canUsesBeRemoved(
     const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return true;
 }
 
 DeletionKind LLVM::LifetimeStartOp::removeBlockingUses(
-    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, RewriterBase &rewriter) {
   return DeletionKind::Delete;
 }
 
 bool LLVM::LifetimeEndOp::canUsesBeRemoved(
     const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return true;
 }
 
 DeletionKind LLVM::LifetimeEndOp::removeBlockingUses(
-    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, RewriterBase &rewriter) {
   return DeletionKind::Delete;
 }
 
 bool LLVM::DbgDeclareOp::canUsesBeRemoved(
     const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return true;
 }
 
 DeletionKind LLVM::DbgDeclareOp::removeBlockingUses(
-    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, RewriterBase &rewriter) {
   return DeletionKind::Delete;
 }
 
 static bool hasAllZeroIndices(LLVM::GEPOp gepOp) {
   return llvm::all_of(gepOp.getIndices(), [](auto index) {
     auto indexAttr = index.template dyn_cast<IntegerAttr>();
     return indexAttr && indexAttr.getValue() == 0;
   });
 }
 
 //===----------------------------------------------------------------------===//
 // Interfaces for GEPOp
 //===----------------------------------------------------------------------===//
 
 bool LLVM::GEPOp::canUsesBeRemoved(
     const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   // GEP can be removed as long as it is a no-op and its users can be removed.
   if (!hasAllZeroIndices(*this))
     return false;
   return forwardToUsers(*this, newBlockingUses);
 }
 
 DeletionKind LLVM::GEPOp::removeBlockingUses(
-    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, RewriterBase &rewriter) {
   return DeletionKind::Delete;
 }

mlir/lib/Dialect/MemRef/IR/MemRefMem2Reg.cpp

SmallVector<MemorySlot> memref::AllocaOp::getPromotableSlots() {
   // Make sure the memref contains only a single element.
   if (any_of(type.getShape(), [](uint64_t dim) { return dim != 1; }))
     return {};
 
   return {MemorySlot{getResult(), type.getElementType()}};
 }
 
 Value memref::AllocaOp::getDefaultValue(const MemorySlot &slot,
-                                        OpBuilder &builder) {
+                                        RewriterBase &rewriter) {
   assert(isSupportedElementType(slot.elemType));
   // TODO: support more types.
   return TypeSwitch<Type, Value>(slot.elemType)
       .Case([&](MemRefType t) {
-        return builder.create<memref::AllocaOp>(getLoc(), t);
+        return rewriter.create<memref::AllocaOp>(getLoc(), t);
       })
       .Default([&](Type t) {
-        return builder.create<arith::ConstantOp>(getLoc(), t,
-                                                 builder.getZeroAttr(t));
+        return rewriter.create<arith::ConstantOp>(getLoc(), t,
+                                                  rewriter.getZeroAttr(t));
       });
 }
 
 void memref::AllocaOp::handlePromotionComplete(const MemorySlot &slot,
-                                               Value defaultValue) {
+                                               Value defaultValue,
+                                               RewriterBase &rewriter) {
   if (defaultValue.use_empty())
-    defaultValue.getDefiningOp()->erase();
-  erase();
+    rewriter.eraseOp(defaultValue.getDefiningOp());
+  rewriter.eraseOp(*this);
 }
 
 void memref::AllocaOp::handleBlockArgument(const MemorySlot &slot,
                                            BlockArgument argument,
-                                           OpBuilder &builder) {}
+                                           RewriterBase &rewriter) {}
 
 //===----------------------------------------------------------------------===//
 //  LoadOp/StoreOp interfaces
 //===----------------------------------------------------------------------===//
 
 bool memref::LoadOp::loadsFrom(const MemorySlot &slot) {
   return getMemRef() == slot.ptr;
 }
 
 Value memref::LoadOp::getStored(const MemorySlot &slot) { return {}; }
 
 bool memref::LoadOp::canUsesBeRemoved(
     const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   if (blockingUses.size() != 1)
     return false;
   Value blockingUse = (*blockingUses.begin())->get();
   return blockingUse == slot.ptr && getMemRef() == slot.ptr &&
          getResult().getType() == slot.elemType;
 }
 
 DeletionKind memref::LoadOp::removeBlockingUses(
     const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder, Value reachingDefinition) {
+    RewriterBase &rewriter, Value reachingDefinition) {
   // `canUsesBeRemoved` checked this blocking use must be the loaded slot
   // pointer.
-  getResult().replaceAllUsesWith(reachingDefinition);
+  rewriter.replaceAllUsesWith(getResult(), reachingDefinition);
   return DeletionKind::Delete;
 }
 
 bool memref::StoreOp::loadsFrom(const MemorySlot &slot) { return false; }
 
 Value memref::StoreOp::getStored(const MemorySlot &slot) {
   if (getMemRef() != slot.ptr)
     return {};
   return getValue();
 }
 
 bool memref::StoreOp::canUsesBeRemoved(
     const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   if (blockingUses.size() != 1)
     return false;
   Value blockingUse = (*blockingUses.begin())->get();
   return blockingUse == slot.ptr && getMemRef() == slot.ptr &&
          getValue() != slot.ptr && getValue().getType() == slot.elemType;
 }
 
 DeletionKind memref::StoreOp::removeBlockingUses(
     const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder, Value reachingDefinition) {
+    RewriterBase &rewriter, Value reachingDefinition) {
   return DeletionKind::Delete;
 }

mlir/lib/Transforms/Mem2Reg.cpp

 //===- Mem2Reg.cpp - Promotes memory slots into values ----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/Mem2Reg.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dominance.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/GenericIteratedDominanceFrontier.h"
 

 ///
 /// For further reading, chapter three of SSA-based Compiler Design [1]
 /// showcases SSA construction, where mem2reg is an adaptation of the same
 /// process.
 ///
 /// [1]: Rastello F. & Bouchez Tichadou F., SSA-based Compiler Design (2022),
 ///      Springer.
 
+namespace {
+
+/// Information computed during promotion analysis used to perform actual
+/// promotion.
+struct MemorySlotPromotionInfo {
+  /// Blocks for which at least two definitions of the slot values clash.
+  SmallPtrSet<Block *, 8> mergePoints;
+  /// Contains, for each operation, which uses must be eliminated by promotion.
+  /// This is a DAG structure because if an operation must eliminate some of
+  /// its uses, it is because the defining ops of the blocking uses requested
+  /// it. The defining ops therefore must also have blocking uses or be the
+  /// starting point of the bloccking uses.
+  DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> userToBlockingUses;
+};
+
+/// Computes information for basic slot promotion. This will check that direct
+/// slot promotion can be performed, and provide the information to execute the
+/// promotion. This does not mutate IR.
+class MemorySlotPromotionAnalyzer {
+public:
+  MemorySlotPromotionAnalyzer(MemorySlot slot, DominanceInfo &dominance)
+      : slot(slot), dominance(dominance) {}
+
+  /// Computes the information for slot promotion if promotion is possible,
+  /// returns nothing otherwise.
+  std::optional<MemorySlotPromotionInfo> computeInfo();
+
+private:
+  /// Computes the transitive uses of the slot that block promotion. This finds
+  /// uses that would block the promotion, checks that the operation has a
+  /// solution to remove the blocking use, and potentially forwards the analysis
+  /// if the operation needs further blocking uses resolved to resolve its own
+  /// uses (typically, removing its users because it will delete itself to
+  /// resolve its own blocking uses). This will fail if one of the transitive
+  /// users cannot remove a requested use, and should prevent promotion.
+  LogicalResult computeBlockingUses(
+      DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> &userToBlockingUses);
+
+  /// Computes in which blocks the value stored in the slot is actually used,
+  /// meaning blocks leading to a load. This method uses `definingBlocks`, the
+  /// set of blocks containing a store to the slot (defining the value of the
+  /// slot).
+  SmallPtrSet<Block *, 16>
+  computeSlotLiveIn(SmallPtrSetImpl<Block *> &definingBlocks);
+
+  /// Computes the points in which multiple re-definitions of the slot's value
+  /// (stores) may conflict.
+  void computeMergePoints(SmallPtrSetImpl<Block *> &mergePoints);
+
+  /// Ensures predecessors of merge points can properly provide their current
+  /// definition of the value stored in the slot to the merge point. This can
+  /// notably be an issue if the terminator used does not have the ability to
+  /// forward values through block operands.
+  bool areMergePointsUsable(SmallPtrSetImpl<Block *> &mergePoints);
+
+  MemorySlot slot;
+  DominanceInfo &dominance;
+};
+
+/// The MemorySlotPromoter handles the state of promoting a memory slot. It
+/// wraps a slot and its associated allocator. This will perform the mutation of
+/// IR.
+class MemorySlotPromoter {
+public:
+  MemorySlotPromoter(MemorySlot slot, PromotableAllocationOpInterface allocator,
+                     RewriterBase &rewriter, DominanceInfo &dominance,
+                     MemorySlotPromotionInfo info,
+                     const Mem2RegStatistics &statistics);
+
+  /// Actually promotes the slot by mutating IR. Promoting a slot DOES
+  /// invalidate the MemorySlotPromotionInfo of other slots. Preparation of
+  /// promotion info should NOT be performed in batches.
+  void promoteSlot();
+
+private:
+  /// Computes the reaching definition for all the operations that require
+  /// promotion. `reachingDef` is the value the slot should contain at the
+  /// beginning of the block. This method returns the reached definition at the
+  /// end of the block.
+  Value computeReachingDefInBlock(Block *block, Value reachingDef);
+
+  /// Computes the reaching definition for all the operations that require
+  /// promotion. `reachingDef` corresponds to the initial value the
+  /// slot will contain before any write, typically a poison value.
+  void computeReachingDefInRegion(Region *region, Value reachingDef);
+
+  /// Removes the blocking uses of the slot, in topological order.
+  void removeBlockingUses();
+
+  /// Lazily-constructed default value representing the content of the slot when
+  /// no store has been executed. This function may mutate IR.
+  Value getLazyDefaultValue();
+
+  MemorySlot slot;
+  PromotableAllocationOpInterface allocator;
+  RewriterBase &rewriter;
+  /// Potentially non-initialized default value. Use `getLazyDefaultValue` to
+  /// initialize it on demand.
+  Value defaultValue;
+  /// Contains the reaching definition at this operation. Reaching definitions
+  /// are only computed for promotable memory operations with blocking uses.
+  DenseMap<PromotableMemOpInterface, Value> reachingDefs;
+  DominanceInfo &dominance;
+  MemorySlotPromotionInfo info;
+  const Mem2RegStatistics &statistics;
+};
+
+} // namespace
+
 MemorySlotPromoter::MemorySlotPromoter(
     MemorySlot slot, PromotableAllocationOpInterface allocator,
-    OpBuilder &builder, DominanceInfo &dominance, MemorySlotPromotionInfo info)
-    : slot(slot), allocator(allocator), builder(builder), dominance(dominance),
-      info(std::move(info)) {
+    RewriterBase &rewriter, DominanceInfo &dominance,
+    MemorySlotPromotionInfo info, const Mem2RegStatistics &statistics)
+    : slot(slot), allocator(allocator), rewriter(rewriter),
+      dominance(dominance), info(std::move(info)), statistics(statistics) {
 #ifndef NDEBUG
   auto isResultOrNewBlockArgument = [&]() {
     if (BlockArgument arg = slot.ptr.dyn_cast<BlockArgument>())
       return arg.getOwner()->getParentOp() == allocator;
     return slot.ptr.getDefiningOp() == allocator;
   };
 
   assert(isResultOrNewBlockArgument() &&
          "a slot must be a result of the allocator or an argument of the child "
          "regions of the allocator");
 #endif // NDEBUG
 }
 
 Value MemorySlotPromoter::getLazyDefaultValue() {
   if (defaultValue)
     return defaultValue;
 
-  OpBuilder::InsertionGuard guard(builder);
-  builder.setInsertionPointToStart(slot.ptr.getParentBlock());
-  return defaultValue = allocator.getDefaultValue(slot, builder);
+  RewriterBase::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointToStart(slot.ptr.getParentBlock());
+  return defaultValue = allocator.getDefaultValue(slot, rewriter);
 }
 
 LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
     DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> &userToBlockingUses) {
   // The promotion of an operation may require the promotion of further
   // operations (typically, removing operations that use an operation that must
   // delete itself). We thus need to start from the use of the slot pointer and
   // propagate further requests through the forward slice.
void MemorySlotPromoter::computeReachingDefInRegion(Region *region,
 
   dfsStack.emplace_back<DfsJob>(
       {domTree.getNode(&region->front()), reachingDef});
 
   while (!dfsStack.empty()) {
     DfsJob job = dfsStack.pop_back_val();
     Block *block = job.block->getBlock();
 
     if (info.mergePoints.contains(block)) {
-      BlockArgument blockArgument =
-          block->addArgument(slot.elemType, slot.ptr.getLoc());
-      builder.setInsertionPointToStart(block);
-      allocator.handleBlockArgument(slot, blockArgument, builder);
+      // If the block is a merge point, we need to add a block argument to hold
+      // the selected reaching definition. This has to be a bit complicated
+      // because of RewriterBase limitations: we need to create a new block with
+      // the extra block argument, move the content of the block to the new
+      // block, and replace the block with the new block in the merge point set.
+      SmallVector<Type> argTypes;
+      SmallVector<Location> argLocs;
+      for (BlockArgument arg : block->getArguments()) {
+        argTypes.push_back(arg.getType());
+        argLocs.push_back(arg.getLoc());
+      }
+      argTypes.push_back(slot.elemType);
+      argLocs.push_back(slot.ptr.getLoc());
+      Block *newBlock = rewriter.createBlock(block, argTypes, argLocs);
+
+      info.mergePoints.erase(block);
+      info.mergePoints.insert(newBlock);
+
+      rewriter.replaceAllUsesWith(block, newBlock);
+      rewriter.mergeBlocks(block, newBlock,
+                           newBlock->getArguments().drop_back());
+
+      block = newBlock;
+
+      BlockArgument blockArgument = block->getArguments().back();
+      rewriter.setInsertionPointToStart(block);
+      allocator.handleBlockArgument(slot, blockArgument, rewriter);
       job.reachingDef = blockArgument;
+
+      if (statistics.newBlockArgumentAmount)
+        (*statistics.newBlockArgumentAmount)++;
     }
 
     job.reachingDef = computeReachingDefInBlock(block, job.reachingDef);
 
     if (auto terminator = dyn_cast<BranchOpInterface>(block->getTerminator())) {
       for (BlockOperand &blockOperand : terminator->getBlockOperands()) {
         if (info.mergePoints.contains(blockOperand.get())) {
           if (!job.reachingDef)
             job.reachingDef = getLazyDefaultValue();
+          rewriter.updateRootInPlace(terminator, [&]() {
             terminator.getSuccessorOperands(blockOperand.getOperandNumber())
                 .append(job.reachingDef);
+          });
         }
       }
     }
 
     for (auto *child : job.block->children())
       dfsStack.emplace_back<DfsJob>({child, job.reachingDef});
   }
 }
void MemorySlotPromoter::removeBlockingUses() {
   for (Operation *toPromote : llvm::reverse(sortedUsersToRemoveUses)) {
     if (auto toPromoteMemOp = dyn_cast<PromotableMemOpInterface>(toPromote)) {
       Value reachingDef = reachingDefs.lookup(toPromoteMemOp);
       // If no reaching definition is known, this use is outside the reach of
       // the slot. The default value should thus be used.
       if (!reachingDef)
         reachingDef = getLazyDefaultValue();
 
-      builder.setInsertionPointAfter(toPromote);
+      rewriter.setInsertionPointAfter(toPromote);
       if (toPromoteMemOp.removeBlockingUses(
-              slot, info.userToBlockingUses[toPromote], builder, reachingDef) ==
-          DeletionKind::Delete)
+              slot, info.userToBlockingUses[toPromote], rewriter,
+              reachingDef) == DeletionKind::Delete)
         toErase.push_back(toPromote);
 
       continue;
     }
 
     auto toPromoteBasic = cast<PromotableOpInterface>(toPromote);
-    builder.setInsertionPointAfter(toPromote);
+    rewriter.setInsertionPointAfter(toPromote);
     if (toPromoteBasic.removeBlockingUses(info.userToBlockingUses[toPromote],
-                                          builder) == DeletionKind::Delete)
+                                          rewriter) == DeletionKind::Delete)
       toErase.push_back(toPromote);
   }
 
   for (Operation *toEraseOp : toErase)
-    toEraseOp->erase();
+    rewriter.eraseOp(toEraseOp);
 
   assert(slot.ptr.use_empty() &&
          "after promotion, the slot pointer should not be used anymore");
 }
 
 void MemorySlotPromoter::promoteSlot() {
   computeReachingDefInRegion(slot.ptr.getParentRegion(), {});
 
   // Now that reaching definitions are known, remove all users.
   removeBlockingUses();
 
   // Update terminators in dead branches to forward default if they are
   // succeeded by a merge points.
   for (Block *mergePoint : info.mergePoints) {
     for (BlockOperand &use : mergePoint->getUses()) {
       auto user = cast<BranchOpInterface>(use.getOwner());
       SuccessorOperands succOperands =
           user.getSuccessorOperands(use.getOperandNumber());
       assert(succOperands.size() == mergePoint->getNumArguments() ||
              succOperands.size() + 1 == mergePoint->getNumArguments());
       if (succOperands.size() + 1 == mergePoint->getNumArguments())
-        succOperands.append(getLazyDefaultValue());
+        rewriter.updateRootInPlace(
+            user, [&]() { succOperands.append(getLazyDefaultValue()); });
     }
   }
 
   LLVM_DEBUG(llvm::dbgs() << "[mem2reg] Promoted memory slot: " << slot.ptr
                           << "\n");
 
-  allocator.handlePromotionComplete(slot, defaultValue);
+  if (statistics.promotedAmount)
+    (*statistics.promotedAmount)++;
+
+  allocator.handlePromotionComplete(slot, defaultValue, rewriter);
 }
 
 LogicalResult mlir::tryToPromoteMemorySlots(
-    ArrayRef<PromotableAllocationOpInterface> allocators, OpBuilder &builder,
-    DominanceInfo &dominance) {
-  // Actual promotion may invalidate the dominance analysis, so slot promotion
-  // is prepated in batches.
-  SmallVector<MemorySlotPromoter> toPromote;
+    ArrayRef<PromotableAllocationOpInterface> allocators,
+    RewriterBase &rewriter, const Mem2RegStatistics &statistics) {
+  DominanceInfo dominance;
+
+  bool promotedAny = false;
+
   for (PromotableAllocationOpInterface allocator : allocators) {
     for (MemorySlot slot : allocator.getPromotableSlots()) {
       if (slot.ptr.use_empty())
         continue;
 
+      DominanceInfo dominance;
       MemorySlotPromotionAnalyzer analyzer(slot, dominance);
       std::optional<MemorySlotPromotionInfo> info = analyzer.computeInfo();
-      if (info)
-        toPromote.emplace_back(slot, allocator, builder, dominance,
-                               std::move(*info));
+      if (info) {
+        MemorySlotPromoter(slot, allocator, rewriter, dominance,
+                           std::move(*info), statistics)
+            .promoteSlot();
+        promotedAny = true;
+      }
     }
   }
 
-  for (MemorySlotPromoter &promoter : toPromote)
-    promoter.promoteSlot();
-
-  return success(!toPromote.empty());
+  return success(promotedAny);
 }
 
-LogicalResult Mem2RegPattern::matchAndRewrite(Operation *op,
-                                              PatternRewriter &rewriter) const {
+LogicalResult
+Mem2RegPattern::matchAndRewrite(PromotableAllocationOpInterface allocator,
   hasBoundedRewriteRecursion();
-
-  if (op->getNumRegions() == 0)
-    return failure();
-
-  DominanceInfo dominance;
-
-  SmallVector<PromotableAllocationOpInterface> allocators;
-  // Build a list of allocators to attempt to promote the slots of.
-  for (Region &region : op->getRegions())
-    for (auto allocator : region.getOps<PromotableAllocationOpInterface>())
-      allocators.emplace_back(allocator);
-
-  // Because pattern rewriters are normally not expressive enough to support a
-  // transformation like mem2reg, this uses an escape hatch to mark modified
-  // operations manually and operate outside of its context.
-  rewriter.startRootUpdate(op);
-
-  OpBuilder builder(rewriter.getContext());
-
-  if (failed(tryToPromoteMemorySlots(allocators, builder, dominance))) {
-    rewriter.cancelRootUpdate(op);
-    return failure();
-  }
-
-  rewriter.finalizeRootUpdate(op);
-  return success();
+  return tryToPromoteMemorySlots({allocator}, rewriter, statistics);
 }
 
 namespace {
 
 struct Mem2Reg : impl::Mem2RegBase<Mem2Reg> {
+  using impl::Mem2RegBase<Mem2Reg>::Mem2RegBase;
+
   void runOnOperation() override {
     Operation *scopeOp = getOperation();
-    bool changed = false;
 
     RewritePatternSet rewritePatterns(&getContext());
-    rewritePatterns.add<Mem2RegPattern>(&getContext());
+    rewritePatterns.add<Mem2RegPattern>(
+        &getContext(),
+        Mem2RegStatistics{&promotedAmount, &newBlockArgumentAmount});
     FrozenRewritePatternSet frozen(std::move(rewritePatterns));
-    (void)applyOpPatternsAndFold({scopeOp}, frozen, GreedyRewriteConfig(),
-                                 &changed);
 
-    if (!changed)
-      markAllAnalysesPreserved();
+    GreedyRewriteConfig config;
+    config.enableRegionSimplification = enableRegionSimplification;
+
+    for (Region &region : scopeOp->getRegions())
+      if (failed(applyPatternsAndFoldGreedily(region, frozen, config)))
+        signalPassFailure();
   }
 };
 
 } // namespace

mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir

-// RUN: mlir-opt %s --pass-pipeline='builtin.module(llvm.func(mem2reg))' | FileCheck %s
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(llvm.func(mem2reg{region-simplify=false}))' | FileCheck %s
 
 llvm.func @use(i64)
 llvm.func @use_ptr(!llvm.ptr)
 
 #di_basic_type = #llvm.di_basic_type<tag = DW_TAG_base_type, name = "ptr sized type", sizeInBits = 64>
 #di_file = #llvm.di_file<"test.ll" in "">
 #di_compile_unit = #llvm.di_compile_unit<sourceLanguage = DW_LANG_C_plus_plus_14, file = #di_file, producer = "clang", isOptimized = false, emissionKind = Full>
 #di_subprogram = #llvm.di_subprogram<compileUnit = #di_compile_unit, scope = #di_file, name = "blah", linkageName = "blah", file = #di_file, line = 7, subprogramFlags = Definition>

mlir/test/Dialect/LLVMIR/mem2reg.mlir

-// RUN: mlir-opt %s --pass-pipeline="builtin.module(llvm.func(mem2reg))" --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(llvm.func(mem2reg{region-simplify=false}))" --split-input-file | FileCheck %s
 
 // CHECK-LABEL: llvm.func @default_value
 llvm.func @default_value() -> i32 {
   // CHECK: %[[UNDEF:.*]] = llvm.mlir.undef : i32
   %0 = llvm.mlir.constant(1 : i32) : i32
   %1 = llvm.alloca %0 x i32 {alignment = 4 : i64} : (i32) -> !llvm.ptr
   %2 = llvm.load %1 {alignment = 4 : i64} : !llvm.ptr -> i32
   // CHECK: llvm.return %[[UNDEF]] : i32

mlir/test/Dialect/MemRef/mem2reg.mlir

-// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(mem2reg))' --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --pass-pipeline='builtin.module(func.func(mem2reg{region-simplify=false}))' --split-input-file | FileCheck %s
 
 // CHECK-LABEL: func.func @basic
 func.func @basic() -> i32 {
   // CHECK-NOT: = memref.alloca
   // CHECK: %[[RES:.*]] = arith.constant 5 : i32
   // CHECK-NOT: = memref.alloca
   %0 = arith.constant 5 : i32
   %1 = memref.alloca() : memref<i32>
func.func @deny_store_of_alloca(%arg: memref<memref<i32>>) -> i32 {
   // CHECK: return %[[RES]] : i32
   return %2 : i32
 }
 
 // -----
 
 // CHECK-LABEL: func.func @promotable_nonpromotable_intertwined
 func.func @promotable_nonpromotable_intertwined() -> i32 {
-  // CHECK: %[[VAL:.*]] = arith.constant 5 : i32
-  %0 = arith.constant 5 : i32
   // CHECK: %[[NON_PROMOTED:.*]] = memref.alloca() : memref<i32>
-  %1 = memref.alloca() : memref<i32>
+  %0 = memref.alloca() : memref<i32>
   // CHECK-NOT: = memref.alloca() : memref<memref<i32>>
-  %2 = memref.alloca() : memref<memref<i32>>
-  memref.store %1, %2[] : memref<memref<i32>>
-  %3 = memref.load %2[] : memref<memref<i32>>
+  %1 = memref.alloca() : memref<memref<i32>>
+  memref.store %0, %1[] : memref<memref<i32>>
+  %2 = memref.load %1[] : memref<memref<i32>>
   // CHECK: call @use(%[[NON_PROMOTED]])
-  call @use(%1) : (memref<i32>) -> ()
+  call @use(%0) : (memref<i32>) -> ()
   // CHECK: %[[RES:.*]] = memref.load %[[NON_PROMOTED]][]
-  %4 = memref.load %1[] : memref<i32>
+  %3 = memref.load %0[] : memref<i32>
   // CHECK: return %[[RES]] : i32
-  return %4 : i32
+  return %3 : i32
 }
 
 func.func @use(%arg: memref<i32>) { return }