[mlir][mem2reg] Add mem2reg rewrite pattern.

[Moxinilian]

Summary

This revision introduces the ability to invoke mem2reg as a rewrite pattern. This also modified the canonical mem2reg pass to use the rewrite pattern approach.

Depends on D149825

[gysit]

Is there logic in the existing mem2reg implementation that could possibly be simplified now that we are using a pattern. For example, applyOpPatternsAndFold should cleanup dead blocks by itself AFAIK. Or is the goal to make tryToPromoteMemorySlots completely self contained to make sure it can be used in a pass without running SimplifyRegion or similar after its execution.

[Moxinilian]

In D149958#4326202, @gysit wrote:

Is there logic in the existing mem2reg implementation that could possibly be simplified now that we are using a pattern. For example, applyOpPatternsAndFold should cleanup dead blocks by itself AFAIK. Or is the goal to make tryToPromoteMemorySlots completely self contained to make sure it can be used in a pass without running SimplifyRegion or similar after its execution.

The intent was the latter yes. Although most of the cleanup in the current implementation is for correctness (except maybe the cleanup in the interfaces): if dead blocks are not properly rewired, the IR structure becomes invalid. So it is definitely necessary if canonicalization is not ran before applying the pattern, which feels like an unnecessary source of fragility.

[gysit]

The intent was the latter yes. Although most of the cleanup in the current implementation is for correctness (except maybe the cleanup in the interfaces): if dead blocks are not properly rewired, the IR structure becomes invalid. So it is definitely necessary if canonicalization is not ran before applying the pattern, which feels like an unnecessary source of fragility.

Ok, I agree that tryToPromoteMemorySlots definitely should not produce invalid IR. Theoretically we could fail if canonicalization did not run but I think it is better to keep the code and have a robust version of mem2reg then!

[gysit]

mlir/lib/Transforms/Mem2Reg.cpp

465

I would add a check if the operation has no region right at the beginning of the patter, e.g.:

if (op->getNumRegions() == 0)
 return failure();

That way the patterns fails early in most cases. Maybe there is also a way to use an OpTrait or OpInterfaceRewritePattern to make sure we are only matching on Operations with a region?

[Dinistro]

mlir/lib/Transforms/Mem2Reg.cpp

466

NIT: could you use region.getOps<PromotableAllocationOpInterface>() here to simplify this?

471–474

Just for my understand: You use Operation::erase() inside the tryToPromoteMermorySlots function and this is the reason why you need a separate builder here, right?

[Moxinilian]

Address comments.

mlir/lib/Transforms/Mem2Reg.cpp

471–474

It's a bit more than that. Not only is Operation::erase used, but operations are modified in place and block argument structure gets modified, this is not supported by RewriterBase.

[gysit]

LGTM modulo one minor comment regarding the debug printout.

mlir/lib/Transforms/Mem2Reg.cpp

425–426

Usually there is a:

#define DEBUG_TYPE "mem2reg"

Add the top of the file. After adding it you should be able to simplify the debug print out to:

LLVM_DEBUG(llvm::dbgs() << "[mem2reg] Promoted memory slot: " << slot.ptr << "\n");

[Moxinilian]

Address debug print comment.

[Moxinilian]

Rebase to latest main.

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"
 
 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.
private:
   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;
 };
 
+/// Pattern applying mem2reg to the regions of the operations on which it
+/// matches.
+class Mem2RegPattern : public RewritePattern {
+public:
+  using RewritePattern::RewritePattern;
+
+  Mem2RegPattern(MLIRContext *ctx, PatternBenefit benefit = 1)
+      : RewritePattern(MatchAnyOpTypeTag(), benefit, ctx) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override;
+};
+
 /// 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);
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_MEM2REG_H

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/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"
 
 namespace mlir {
 #define GEN_PASS_DEF_MEM2REG
 #include "mlir/Transforms/Passes.h.inc"
for (BlockOperand &use : mergePoint->getUses()) {
           user.getSuccessorOperands(use.getOperandNumber());
       assert(succOperands.size() == mergePoint->getNumArguments() ||
              succOperands.size() + 1 == mergePoint->getNumArguments());
       if (succOperands.size() + 1 == mergePoint->getNumArguments())
         succOperands.append(getLazyDefaultValue());
     }
   }
 
+  DEBUG_WITH_TYPE("mem2reg", llvm::dbgs() << "[mem2reg] Promoted memory slot: "
+                                          << slot.ptr << "\n");
+
   allocator.handlePromotionComplete(slot, defaultValue);
 }
 
 LogicalResult mlir::tryToPromoteMemorySlots(
     ArrayRef<PromotableAllocationOpInterface> allocators, OpBuilder &builder,
     DominanceInfo &dominance) {
   // Actual promotion may invalidate the dominance analysis, so slot promotion
   // is prepated in batches.
LogicalResult mlir::tryToPromoteMemorySlots(
   }
 
   for (MemorySlotPromoter &promoter : toPromote)
     promoter.promoteSlot();
 
   return success(!toPromote.empty());
 }
 
-namespace {
-
-struct Mem2Reg : impl::Mem2RegBase<Mem2Reg> {
-  void runOnOperation() override {
-    Operation *scopeOp = getOperation();
-    bool changed = false;
-
-    for (Region &region : scopeOp->getRegions()) {
-      if (region.getBlocks().empty())
-        continue;
+LogicalResult Mem2RegPattern::matchAndRewrite(Operation *op,
+                                              PatternRewriter &rewriter) const {
+  hasBoundedRewriteRecursion();
 
-      OpBuilder builder(&region.front(), region.front().begin());
+  if (op->getNumRegions() == 0)
+    return failure();
 
-      // Promoting a slot can allow for further promotion of other slots,
-      // promotion is tried until no promotion succeeds.
-      while (true) {
-        DominanceInfo &dominance = getAnalysis<DominanceInfo>();
+  DominanceInfo dominance;
 
-        SmallVector<PromotableAllocationOpInterface> allocators;
-        // Build a list of allocators to attempt to promote the slots of.
-        for (Block &block : region)
-          for (Operation &op : block.getOperations())
-            if (auto allocator = dyn_cast<PromotableAllocationOpInterface>(op))
-              allocators.emplace_back(allocator);
+  for (Region &region : op->getRegions())
+    for (auto allocator : region.getOps<PromotableAllocationOpInterface>())
 
-        // Attempt promoting until no promotion succeeds.
-        if (failed(tryToPromoteMemorySlots(allocators, builder, dominance)))
-          break;
+  // 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());
 
-        changed = true;
-        getAnalysisManager().invalidate({});
-      }
-    }
+  if (failed(tryToPromoteMemorySlots(allocators, builder, dominance))) {
+    rewriter.cancelRootUpdate(op);
+    return failure();
+
+  rewriter.finalizeRootUpdate(op);
+  return success();
 
+namespace {
+
+struct Mem2Reg : impl::Mem2RegBase<Mem2Reg> {
+  void runOnOperation() override {
+    Operation *scopeOp = getOperation();
+    bool changed = false;
+
+    RewritePatternSet rewritePatterns(&getContext());
+    rewritePatterns.add<Mem2RegPattern>(&getContext());
+    FrozenRewritePatternSet frozen(std::move(rewritePatterns));
+    (void)applyOpPatternsAndFold({scopeOp}, frozen, GreedyRewriteConfig(),
+                                 &changed);
+
     if (!changed)
       markAllAnalysesPreserved();
   }
 };
 
 } // namespace