[mlir][mem2reg] Expose algorithm internals.

[Moxinilian]

Summary

This patch refactors the Mem2Reg infrastructure. It decouples
analysis from promotion, allowing for more control over the execution of
the logic. It also adjusts the interfaces to be less coupled to mem2reg
and more general. This will be useful for an upcoming revision
introducing generic SROA.

[Moxinilian]

Rebase to latest main.

[gysit]

mlir/include/mlir/Transforms/Mem2Reg.h

25

nit: Can you put a comma after the first must eliminate or reformulate a bit to simplify the sentence?

40

Optional is deprecated and should be replaced by std::optional AFAIK.

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

20

Are you sure all this additional includes are needed? I have the feeling some of them can be eliminated?

mlir/lib/Transforms/Mem2Reg.cpp

298

nit: I would go for return info;. I would expect the compiler can do copy elision here (even if you use std::optional). The main argument is that it is more readable though.

441

nit: Once you use std::optional please do not use has_value() and value(). Instead just use the optional directly in the condition and use * to get its value.

[Moxinilian]

Address review comments.

[Moxinilian]

mlir/include/mlir/Transforms/Mem2Reg.h

25

I reformulated using words that require less context so it should be clearer.

[gysit]

LGTM

[Dinistro]

mlir/include/mlir/Interfaces/CMakeLists.txt

19

What's the reason to make a hand-rolled tablegen target here?

[Moxinilian]

mlir/include/mlir/Interfaces/CMakeLists.txt

19

This is because it was easier for me to extract this change out of my SROA cchange, and I forgot to manually revert it. The quickly upcoming SROA revision will make use of this by generating type interfaces, so I can also just leave it as-is as it will be used shortly. Please tell me what you prefer.

[Dinistro]

mlir/include/mlir/Interfaces/CMakeLists.txt

19

Technically, this change is not necessary. On the other hand, this will change anyway soonish, so don't bother with reverting this part.

[Moxinilian]

Rebased to latest main.

[Moxinilian]

Reopen due to accidental revert.

[Moxinilian]

Regenerate commit.

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

 #ifndef LLVM_INTRINSIC_OPS
 #define LLVM_INTRINSIC_OPS
 
 include "mlir/IR/OpBase.td"
 include "mlir/Dialect/LLVMIR/LLVMAttrDefs.td"
 include "mlir/Dialect/LLVMIR/LLVMEnums.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
-include "mlir/Interfaces/Mem2RegInterfaces.td"
+include "mlir/Interfaces/MemorySlotInterfaces.td"
 
 // Operations that correspond to LLVM intrinsics. With MLIR operation set being
 // extendable, there is no reason to introduce a hard boundary between "core"
 // operations and intrinsics. However, we systematically prefix them with
 // "intr." to avoid potential name clashes.
 
 class LLVM_UnaryIntrOpBase<string func, Type element,
                            list<Trait> traits = [], bit requiresFastmath = 0> :

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 include "mlir/Dialect/LLVMIR/LLVMEnums.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/FunctionInterfaces.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
-include "mlir/Interfaces/Mem2RegInterfaces.td"
+include "mlir/Interfaces/MemorySlotInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class LLVM_Builder<string builder> {
   string llvmBuilder = builder;
 }
 
 // Base class for LLVM terminator operations.  All terminator operations have
 // zero results and an optional list of successors.

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

 
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
-#include "mlir/Interfaces/Mem2RegInterfaces.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include "mlir/Interfaces/ShapedOpInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include <optional>
 
 namespace mlir {
 
 namespace arith {

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 #define MEMREF_OPS
 
 include "mlir/Dialect/Arith/IR/ArithBase.td"
 include "mlir/Dialect/MemRef/IR/MemRefBase.td"
 include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
-include "mlir/Interfaces/Mem2RegInterfaces.td"
+include "mlir/Interfaces/MemorySlotInterfaces.td"
 include "mlir/Interfaces/ShapedOpInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 
 /// A TypeAttr for memref types.
 def MemRefTypeAttr

mlir/include/mlir/Interfaces/CMakeLists.txt

 add_mlir_interface(CallInterfaces)
 add_mlir_interface(CastInterfaces)
 add_mlir_interface(ControlFlowInterfaces)
 add_mlir_interface(CopyOpInterface)
 add_mlir_interface(DerivedAttributeOpInterface)
 add_mlir_interface(DestinationStyleOpInterface)
 add_mlir_interface(InferIntRangeInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)
-add_mlir_interface(Mem2RegInterfaces)
 add_mlir_interface(ParallelCombiningOpInterface)
 add_mlir_interface(RuntimeVerifiableOpInterface)
 add_mlir_interface(ShapedOpInterfaces)
 add_mlir_interface(SideEffectInterfaces)
 add_mlir_interface(TilingInterface)
 add_mlir_interface(ValueBoundsOpInterface)
 add_mlir_interface(VectorInterfaces)
 add_mlir_interface(ViewLikeInterface)
 
+set(LLVM_TARGET_DEFINITIONS MemorySlotInterfaces.td)
+mlir_tablegen(MemorySlotOpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(MemorySlotOpInterfaces.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRMemorySlotInterfacesIncGen)
+add_dependencies(mlir-generic-headers MLIRMemorySlotInterfacesIncGen)
+
 set(LLVM_TARGET_DEFINITIONS DataLayoutInterfaces.td)
 mlir_tablegen(DataLayoutAttrInterface.h.inc -gen-attr-interface-decls)
 mlir_tablegen(DataLayoutAttrInterface.cpp.inc -gen-attr-interface-defs)
 mlir_tablegen(DataLayoutOpInterface.h.inc -gen-op-interface-decls)
 mlir_tablegen(DataLayoutOpInterface.cpp.inc -gen-op-interface-defs)
 mlir_tablegen(DataLayoutTypeInterface.h.inc -gen-type-interface-decls)
 mlir_tablegen(DataLayoutTypeInterface.cpp.inc -gen-type-interface-defs)
 add_public_tablegen_target(MLIRDataLayoutInterfacesIncGen)

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_MEM2REGINTERFACES_H
-#define MLIR_INTERFACES_MEM2REGINTERFACES_H
+#ifndef MLIR_INTERFACES_MEMORYSLOTINTERFACES_H
+#define MLIR_INTERFACES_MEMORYSLOTINTERFACES_H
 
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/OpDefinition.h"
 
 namespace mlir {
 
 /// Represents a slot in memory. This is generated by an allocating operation
 /// (for example alloca).
enum class DeletionKind {
   /// Keep the operation after promotion.
   Keep,
   /// Delete the operation after promotion.
   Delete,
 };
 
 } // namespace mlir
 
-#include "mlir/Interfaces/Mem2RegInterfaces.h.inc"
+#include "mlir/Interfaces/MemorySlotOpInterfaces.h.inc"
 
-#endif // MLIR_INTERFACES_MEM2REGINTERFACES_H
+#endif // MLIR_INTERFACES_MEMORYSLOTINTERFACES_H

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

-//===-- Mem2RegInterfaces.td - Mem2Reg interfaces ----------*- tablegen -*-===//
+//===-- MemorySlotInterfaces.td - MemorySlot interfaces ----*- tablegen -*-===//
 //
 // 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_MEM2REGINTERFACES
-#define MLIR_INTERFACES_MEM2REGINTERFACES
+#ifndef MLIR_INTERFACES_MEMORYSLOTINTERFACES
+#define MLIR_INTERFACES_MEMORYSLOTINTERFACES
 
 include "mlir/IR/OpBase.td"
 
 def PromotableAllocationOpInterface
     : OpInterface<"PromotableAllocationOpInterface"> {
   let description = [{
     Describes an operation allocating a memory slot that can be promoted into
     SSA values.

 }
 
 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.
let description = [{
     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 the context of promoting `slot`.
+        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.
       }], "bool", "canUsesBeRemoved",
-      (ins "const ::mlir::MemorySlot &":$slot,
-           "const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &":$blockingUses,
+      (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 memory slot anymore. In contrast to `PromotableMemOpInterface`,
-        operations implementing this interface must not need access to the
-        reaching definition of the content of the slot.
+        provided blocking uses anymore. In contrast to
+        `PromotableMemOpInterface`, operations implementing this interface
+        must not need access to the reaching definition of the content of the
+        slot.
 
         During the transformation, *no operation should be deleted*.
         The operation can only schedule its own deletion by returning the
         appropriate `DeletionKind`. The deletion must be legal assuming the
         blocking uses passed through the `newBlockingUses` list in
         `canUseBeRemoved` have been removed.
 
         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.
       }],
       "::mlir::DeletionKind",
       "removeBlockingUses",
-      (ins "const ::mlir::MemorySlot &":$slot,
-           "const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
+      (ins "const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
            "::mlir::OpBuilder &":$builder)
     >,
   ];
 }
 
-#endif // MLIR_INTERFACES_MEM2REGINTERFACES
+#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/Interfaces/Mem2RegInterfaces.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.
+  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;
+};
+
 /// 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/Dialect/LLVMIR/CMakeLists.txt

 add_subdirectory(Transforms)
 
 add_mlir_dialect_library(MLIRLLVMDialect
   IR/FunctionCallUtils.cpp
   IR/LLVMAttrs.cpp
   IR/LLVMDialect.cpp
   IR/LLVMInlining.cpp
   IR/LLVMInterfaces.cpp
-  IR/LLVMMem2Reg.cpp
+  IR/LLVMMemorySlot.cpp
   IR/LLVMTypes.cpp
   IR/LLVMTypeSyntax.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/LLVMIR
 
   DEPENDS
   MLIRLLVMOpsIncGen
add_mlir_dialect_library(MLIRLLVMDialect
   Core
 
   LINK_LIBS PUBLIC
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces
   MLIRInferTypeOpInterface
   MLIRIR
-  MLIRMem2RegInterfaces
+  MLIRMemorySlotInterfaces
   MLIRSideEffectInterfaces
   MLIRSupport
   )
 
 add_mlir_dialect_library(MLIRNVVMDialect
   IR/NVVMDialect.cpp
 
   ADDITIONAL_HEADER_DIRS

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

-//===- LLVMMem2Reg.cpp - Mem2Reg Interfaces ---------------------*- C++ -*-===//
+//===- LLVMMemorySlot.cpp - MemorySlot 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
 //
 //===----------------------------------------------------------------------===//
 //
-// This file implements Mem2Reg-related interfaces for LLVM dialect operations.
+// This file implements MemorySlot-related interfaces for LLVM dialect
+// operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
-#include "mlir/Interfaces/Mem2RegInterfaces.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Interfaces for AllocaOp
 //===----------------------------------------------------------------------===//
 
+static Type getAllocaElementType(LLVM::AllocaOp alloca) {
+  // This will become trivial once non-opaque pointers are gone.
+  return alloca.getElemType().has_value()
+             ? *alloca.getElemType()
+             : alloca.getResult().getType().getElementType();
+}
+
 llvm::SmallVector<MemorySlot> LLVM::AllocaOp::getPromotableSlots() {
   if (!getOperation()->getBlock()->isEntryBlock())
     return {};
 
-  Type elemType =
-      getElemType() ? *getElemType() : getResult().getType().getElementType();
-  return {MemorySlot{getResult(), elemType}};
+  return {MemorySlot{getResult(), getAllocaElementType(*this)}};
 }
 
 Value LLVM::AllocaOp::getDefaultValue(const MemorySlot &slot,
                                       OpBuilder &builder) {
   return builder.create<LLVM::UndefOp>(getLoc(), slot.elemType);
 }
 
 void LLVM::AllocaOp::handleBlockArgument(const MemorySlot &slot,
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 MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return forwardToUsers(*this, newBlockingUses);
 }
 
 DeletionKind LLVM::BitcastOp::removeBlockingUses(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
   return DeletionKind::Delete;
 }
 
 bool LLVM::AddrSpaceCastOp::canUsesBeRemoved(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return forwardToUsers(*this, newBlockingUses);
 }
 
 DeletionKind LLVM::AddrSpaceCastOp::removeBlockingUses(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
   return DeletionKind::Delete;
 }
 
 bool LLVM::LifetimeStartOp::canUsesBeRemoved(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return true;
 }
 
 DeletionKind LLVM::LifetimeStartOp::removeBlockingUses(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
   return DeletionKind::Delete;
 }
 
 bool LLVM::LifetimeEndOp::canUsesBeRemoved(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return true;
 }
 
 DeletionKind LLVM::LifetimeEndOp::removeBlockingUses(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
   return DeletionKind::Delete;
 }
 
 bool LLVM::DbgDeclareOp::canUsesBeRemoved(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses) {
   return true;
 }
 
 DeletionKind LLVM::DbgDeclareOp::removeBlockingUses(
-    const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
   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 MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    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 MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
-    OpBuilder &builder) {
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
   return DeletionKind::Delete;
 }

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

 //===- MemRefMem2Reg.cpp - 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements Mem2Reg-related interfaces for MemRef dialect
 // operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Interfaces/Mem2RegInterfaces.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 //  AllocaOp interfaces
 //===----------------------------------------------------------------------===//
 

mlir/lib/Interfaces/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   CallInterfaces.cpp
   CastInterfaces.cpp
   ControlFlowInterfaces.cpp
   CopyOpInterface.cpp
   DataLayoutInterfaces.cpp
   DerivedAttributeOpInterface.cpp
   DestinationStyleOpInterface.cpp
   InferIntRangeInterface.cpp
   InferTypeOpInterface.cpp
   LoopLikeInterface.cpp
-  Mem2RegInterfaces.cpp
+  MemorySlotInterfaces.cpp
   ParallelCombiningOpInterface.cpp
   RuntimeVerifiableOpInterface.cpp
   ShapedOpInterfaces.cpp
   SideEffectInterfaces.cpp
   TilingInterface.cpp
   ValueBoundsOpInterface.cpp
   VectorInterfaces.cpp
   ViewLikeInterface.cpp

 add_mlir_interface_library(ControlFlowInterfaces)
 add_mlir_interface_library(CopyOpInterface)
 add_mlir_interface_library(DataLayoutInterfaces)
 add_mlir_interface_library(DerivedAttributeOpInterface)
 add_mlir_interface_library(DestinationStyleOpInterface)
 add_mlir_interface_library(InferIntRangeInterface)
 add_mlir_interface_library(InferTypeOpInterface)
 add_mlir_interface_library(LoopLikeInterface)
-add_mlir_interface_library(Mem2RegInterfaces)
+add_mlir_interface_library(MemorySlotInterfaces)
 add_mlir_interface_library(ParallelCombiningOpInterface)
 add_mlir_interface_library(RuntimeVerifiableOpInterface)
 add_mlir_interface_library(ShapedOpInterfaces)
 add_mlir_interface_library(SideEffectInterfaces)
 add_mlir_interface_library(TilingInterface)
 add_mlir_interface_library(VectorInterfaces)
 add_mlir_interface_library(ViewLikeInterface)
 

mlir/lib/Interfaces/MemorySlotInterfaces.cpp

-//===-- Mem2RegInterfaces.cpp - Mem2Reg interfaces --------------*- C++ -*-===//
+//===-- MemorySlotInterfaces.cpp - MemorySlot 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
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Interfaces/Mem2RegInterfaces.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
 
-#include "mlir/Interfaces/Mem2RegInterfaces.cpp.inc"
+#include "mlir/Interfaces/MemorySlotOpInterfaces.cpp.inc"

mlir/lib/Transforms/CMakeLists.txt

add_mlir_library(MLIRTransforms
 
   DEPENDS
   MLIRTransformsPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
   MLIRCopyOpInterface
   MLIRLoopLikeInterface
-  MLIRMem2RegInterfaces
+  MLIRMemorySlotInterfaces
   MLIRPass
   MLIRRuntimeVerifiableOpInterface
   MLIRSideEffectInterfaces
   MLIRSupport
   MLIRTransformUtils
   )

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/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"
 } // namespace mlir
 
 using namespace mlir;

 /// Control flow can create situations where a load could be replaced by
 /// multiple possible stores depending on the control flow path taken. As a
 /// result, this pass must introduce new block arguments in some blocks to
 /// accomodate for the multiple possible definitions. Each predecessor will
 /// populate the block argument with the definition reached at its end. With
 /// this, the value stored can be well defined at block boundaries, allowing
 /// the propagation of replacement through blocks.
 ///
-/// This pass computes this transformation in four main steps:
+/// This pass computes this transformation in four main steps. The two first
+/// steps are performed during an analysis phase that does not mutate IR.
+///
+/// The two steps of the analysis phase are the following:
 /// - A first step computes the list of operations that transitively use the
 /// memory slot we would like to promote. The purpose of this phase is to
 /// identify which uses must be removed to promote the slot, either by rewiring
 /// the user or deleting it. Naturally, direct uses of the slot must be removed.
 /// Sometimes additional uses must also be removed: this is notably the case
 /// when a direct user of the slot cannot rewire its use and must delete itself,
 /// and thus must make its users no longer use it. If any of those uses cannot
 /// be removed by their users in any way, promotion cannot continue: this is
 /// decided at this step.
 /// - A second step computes the list of blocks where a block argument will be
 /// needed ("merge points") without mutating the IR. These blocks are the blocks
 /// leading to a definition clash between two predecessors. Such blocks happen
 /// to be the Iterated Dominance Frontier (IDF) of the set of blocks containing
 /// a store, as they represent the point where a clear defining dominator stops
 /// existing. Computing this information in advance allows making sure the
 /// terminators that will forward values are capable of doing so (inability to
 /// do so aborts promotion at this step).
+///
+/// At this point, promotion is guaranteed to happen, and the mutation phase can
+/// begin with the following steps:
 /// - A third step computes the reaching definition of the memory slot at each
 /// blocking user. This is the core of the mem2reg algorithm, also known as
 /// load-store forwarding. This analyses loads and stores and propagates which
 /// value must be stored in the slot at each blocking user.  This is achieved by
 /// doing a depth-first walk of the dominator tree of the function. This is
 /// sufficient because the reaching definition at the beginning of a block is
 /// either its new block argument if it is a merge block, or the definition
 /// reaching the end of its immediate dominator (parent in the dominator tree).
 /// We can therefore propagate this information down the dominator tree to
 /// proceed with renaming within blocks.
 /// - The final fourth step uses the reaching definition to remove blocking uses
 /// in topological order.
 ///
-/// The two first steps do not mutate IR because promotion can still be aborted
-/// at this point. Once the two last steps are reached, promotion is guaranteed
-/// to succeed, allowing to start mutating IR.
-///
 /// 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 {
-
-/// The SlotPromoter handles the state of promoting a memory slot. It wraps a
-/// slot and its associated allocator, along with analysis results related to
-/// the slot.
-class SlotPromoter {
-public:
-  SlotPromoter(MemorySlot slot, PromotableAllocationOpInterface allocator,
-               OpBuilder &builder, DominanceInfo &dominance);
-
-  /// Prepare data for the promotion of the slot while checking if it can be
-  /// promoted. Succeeds if the slot can be promoted. This method does not
-  /// mutate IR.
-  LogicalResult prepareSlotPromotion();
-
-  /// Actually promotes the slot by mutating IR. This method must only be
-  /// called after a successful call to `SlotPromoter::prepareSlotPromotion`.
-  /// Promoting a slot does not invalidate the preparation of other slots.
-  void promoteSlot();
-
-private:
-  /// This is the first step of the promotion algorithm.
-  /// 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();
-
-  /// 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);
-
-  /// This is the second step of the promotion algorithm.
-  /// Computes the points in which multiple re-definitions of the slot's value
-  /// (stores) may conflict.
-  void computeMergePoints();
-
-  /// 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();
-
-  /// 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);
-
-  /// This is the third step of the promotion algorithm.
-  /// 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);
-
-  /// This is the fourth step of the promotion algorithm.
-  /// 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 `lazyDefaultValue` to
-  /// initialize it on demand.
-  Value defaultValue;
-  /// Blocks where multiple definitions of the slot value clash.
-  SmallPtrSet<Block *, 8> mergePoints;
-  /// Contains, for each operation, which uses must be eliminated by promotion.
-  /// This is a DAG structure because an operation that must eliminate some of
-  /// its uses always comes from a request from an operation that must
-  /// eliminate some of its own uses.
-  DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> userToBlockingUses;
-  /// 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;
-};
-
-} // namespace
-
-SlotPromoter::SlotPromoter(MemorySlot slot,
-                           PromotableAllocationOpInterface allocator,
-                           OpBuilder &builder, DominanceInfo &dominance)
-    : slot(slot), allocator(allocator), builder(builder), dominance(dominance) {
+MemorySlotPromoter::MemorySlotPromoter(
+    MemorySlot slot, PromotableAllocationOpInterface allocator,
+    OpBuilder &builder, DominanceInfo &dominance, MemorySlotPromotionInfo info)
+    : slot(slot), allocator(allocator), builder(builder), dominance(dominance),
+      info(std::move(info)) {
 #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 SlotPromoter::getLazyDefaultValue() {
+Value MemorySlotPromoter::getLazyDefaultValue() {
   if (defaultValue)
     return defaultValue;
 
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointToStart(slot.ptr.getParentBlock());
   return defaultValue = allocator.getDefaultValue(slot, builder);
 }
 
-LogicalResult SlotPromoter::computeBlockingUses() {
+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.
 
   // First insert that all immediate users of the slot pointer must no longer
   // use it.
   for (OpOperand &use : slot.ptr.getUses()) {
     SmallPtrSet<OpOperand *, 4> &blockingUses =
         userToBlockingUses.getOrInsertDefault(use.getOwner());
     blockingUses.insert(&use);
   }
 
   // Then, propagate the requirements for the removal of uses. The
   // topologically-sorted forward slice allows for all blocking uses of an
-  // operation to have been computed before we reach it. Operations are
+  // operation to have been computed before it is reached. Operations are
   // traversed in topological order of their uses, starting from the slot
   // pointer.
   SetVector<Operation *> forwardSlice;
   mlir::getForwardSlice(slot.ptr, &forwardSlice);
   for (Operation *user : forwardSlice) {
     // If the next operation has no blocking uses, everything is fine.
     if (!userToBlockingUses.contains(user))
       continue;
 
     SmallPtrSet<OpOperand *, 4> &blockingUses = userToBlockingUses[user];
 
     SmallVector<OpOperand *> newBlockingUses;
     // If the operation decides it cannot deal with removing the blocking uses,
     // promotion must fail.
     if (auto promotable = dyn_cast<PromotableOpInterface>(user)) {
-      if (!promotable.canUsesBeRemoved(slot, blockingUses, newBlockingUses))
+      if (!promotable.canUsesBeRemoved(blockingUses, newBlockingUses))
         return failure();
     } else if (auto promotable = dyn_cast<PromotableMemOpInterface>(user)) {
       if (!promotable.canUsesBeRemoved(slot, blockingUses, newBlockingUses))
         return failure();
     } else {
       // An operation that has blocking uses must be promoted. If it is not
       // promotable, promotion must fail.
       return failure();
     }
 
     // Then, register any new blocking uses for coming operations.
     for (OpOperand *blockingUse : newBlockingUses) {
       assert(llvm::is_contained(user->getResults(), blockingUse->get()));
 
       SmallPtrSetImpl<OpOperand *> &newUserBlockingUseSet =
           userToBlockingUses.getOrInsertDefault(blockingUse->getOwner());
       newUserBlockingUseSet.insert(blockingUse);
     }
   }
 
   // Because this pass currently only supports analysing the parent region of
-  // the slot pointer, if a promotable memory op that needs promotion is
-  // outside of this region, promotion must fail because it will be impossible
-  // to provide a valid `reachingDef` for it.
+  // the slot pointer, if a promotable memory op that needs promotion is outside
+  // of this region, promotion must fail because it will be impossible to
+  // provide a valid `reachingDef` for it.
   for (auto &[toPromote, _] : userToBlockingUses)
     if (isa<PromotableMemOpInterface>(toPromote) &&
         toPromote->getParentRegion() != slot.ptr.getParentRegion())
       return failure();
 
   return success();
 }
 
-SmallPtrSet<Block *, 16>
-SlotPromoter::computeSlotLiveIn(SmallPtrSetImpl<Block *> &definingBlocks) {
+SmallPtrSet<Block *, 16> MemorySlotPromotionAnalyzer::computeSlotLiveIn(
+    SmallPtrSetImpl<Block *> &definingBlocks) {
   SmallPtrSet<Block *, 16> liveIn;
 
   // The worklist contains blocks in which it is known that the slot value is
   // live-in. The further blocks where this value is live-in will be inferred
   // from these.
   SmallVector<Block *> liveInWorkList;
 
   // Blocks with a load before any other store to the slot are the starting
for (Block *pred : liveInBlock->getPredecessors())
       if (!definingBlocks.contains(pred))
         liveInWorkList.push_back(pred);
   }
 
   return liveIn;
 }
 
 using IDFCalculator = llvm::IDFCalculatorBase<Block, false>;
-void SlotPromoter::computeMergePoints() {
+void MemorySlotPromotionAnalyzer::computeMergePoints(
+    SmallPtrSetImpl<Block *> &mergePoints) {
   if (slot.ptr.getParentRegion()->hasOneBlock())
     return;
 
   IDFCalculator idfCalculator(dominance.getDomTree(slot.ptr.getParentRegion()));
 
   SmallPtrSet<Block *, 16> definingBlocks;
   for (Operation *user : slot.ptr.getUsers())
     if (auto storeOp = dyn_cast<PromotableMemOpInterface>(user))
       if (storeOp.getStored(slot))
         definingBlocks.insert(user->getBlock());
 
   idfCalculator.setDefiningBlocks(definingBlocks);
 
   SmallPtrSet<Block *, 16> liveIn = computeSlotLiveIn(definingBlocks);
   idfCalculator.setLiveInBlocks(liveIn);
 
   SmallVector<Block *> mergePointsVec;
   idfCalculator.calculate(mergePointsVec);
 
   mergePoints.insert(mergePointsVec.begin(), mergePointsVec.end());
 }
 
-bool SlotPromoter::areMergePointsUsable() {
+bool MemorySlotPromotionAnalyzer::areMergePointsUsable(
+    SmallPtrSetImpl<Block *> &mergePoints) {
   for (Block *mergePoint : mergePoints)
     for (Block *pred : mergePoint->getPredecessors())
       if (!isa<BranchOpInterface>(pred->getTerminator()))
         return false;
 
   return true;
 }
 
-Value SlotPromoter::computeReachingDefInBlock(Block *block, Value reachingDef) {
+std::optional<MemorySlotPromotionInfo>
+MemorySlotPromotionAnalyzer::computeInfo() {
+  MemorySlotPromotionInfo info;
+
+  // First, find the set of operations that will need to be changed for the
+  // promotion to happen. These operations need to resolve some of their uses,
+  // either by rewiring them or simply deleting themselves. If any of them
+  // cannot find a way to resolve their blocking uses, we abort the promotion.
+  if (failed(computeBlockingUses(info.userToBlockingUses)))
+    return {};
+
+  // Then, compute blocks in which two or more definitions of the allocated
+  // variable may conflict. These blocks will need a new block argument to
+  // accomodate this.
+  computeMergePoints(info.mergePoints);
+
+  // The slot can be promoted if the block arguments to be created can
+  // actually be populated with values, which may not be possible depending
+  // on their predecessors.
+  if (!areMergePointsUsable(info.mergePoints))
+    return {};
+
+  return info;
+}
+
+Value MemorySlotPromoter::computeReachingDefInBlock(Block *block,
+                                                    Value reachingDef) {
   for (Operation &op : block->getOperations()) {
     if (auto memOp = dyn_cast<PromotableMemOpInterface>(op)) {
-      if (userToBlockingUses.contains(memOp))
+      if (info.userToBlockingUses.contains(memOp))
         reachingDefs.insert({memOp, reachingDef});
 
       if (Value stored = memOp.getStored(slot))
         reachingDef = stored;
     }
   }
 
   return reachingDef;
 }
 
-void SlotPromoter::computeReachingDefInRegion(Region *region,
+void MemorySlotPromoter::computeReachingDefInRegion(Region *region,
                                                     Value reachingDef) {
   if (region->hasOneBlock()) {
     computeReachingDefInBlock(&region->front(), reachingDef);
     return;
   }
 
   struct DfsJob {
     llvm::DomTreeNodeBase<Block> *block;
     Value reachingDef;
   };
 
   SmallVector<DfsJob> dfsStack;
 
   auto &domTree = dominance.getDomTree(slot.ptr.getParentRegion());
 
   dfsStack.emplace_back<DfsJob>(
       {domTree.getNode(&region->front()), reachingDef});
 
   while (!dfsStack.empty()) {
     DfsJob job = dfsStack.pop_back_val();
     Block *block = job.block->getBlock();
 
-    if (mergePoints.contains(block)) {
+    if (info.mergePoints.contains(block)) {
       BlockArgument blockArgument =
           block->addArgument(slot.elemType, slot.ptr.getLoc());
       builder.setInsertionPointToStart(block);
       allocator.handleBlockArgument(slot, blockArgument, builder);
       job.reachingDef = blockArgument;
     }
 
     job.reachingDef = computeReachingDefInBlock(block, job.reachingDef);
 
     if (auto terminator = dyn_cast<BranchOpInterface>(block->getTerminator())) {
       for (BlockOperand &blockOperand : terminator->getBlockOperands()) {
-        if (mergePoints.contains(blockOperand.get())) {
+        if (info.mergePoints.contains(blockOperand.get())) {
           if (!job.reachingDef)
             job.reachingDef = getLazyDefaultValue();
           terminator.getSuccessorOperands(blockOperand.getOperandNumber())
               .append(job.reachingDef);
         }
       }
     }
 
     for (auto *child : job.block->children())
       dfsStack.emplace_back<DfsJob>({child, job.reachingDef});
   }
 }
 
-void SlotPromoter::removeBlockingUses() {
+void MemorySlotPromoter::removeBlockingUses() {
   llvm::SetVector<Operation *> usersToRemoveUses;
-  for (auto &user : llvm::make_first_range(userToBlockingUses))
+  for (auto &user : llvm::make_first_range(info.userToBlockingUses))
     usersToRemoveUses.insert(user);
   SetVector<Operation *> sortedUsersToRemoveUses =
       mlir::topologicalSort(usersToRemoveUses);
 
   llvm::SmallVector<Operation *> toErase;
   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);
-      if (toPromoteMemOp.removeBlockingUses(slot, userToBlockingUses[toPromote],
-                                            builder, reachingDef) ==
+      if (toPromoteMemOp.removeBlockingUses(
+              slot, info.userToBlockingUses[toPromote], builder, reachingDef) ==
           DeletionKind::Delete)
         toErase.push_back(toPromote);
 
       continue;
     }
 
     auto toPromoteBasic = cast<PromotableOpInterface>(toPromote);
     builder.setInsertionPointAfter(toPromote);
-    if (toPromoteBasic.removeBlockingUses(slot, userToBlockingUses[toPromote],
+    if (toPromoteBasic.removeBlockingUses(info.userToBlockingUses[toPromote],
                                           builder) == DeletionKind::Delete)
       toErase.push_back(toPromote);
   }
 
   for (Operation *toEraseOp : toErase)
     toEraseOp->erase();
 
   assert(slot.ptr.use_empty() &&
          "after promotion, the slot pointer should not be used anymore");
 }
 
-void SlotPromoter::promoteSlot() {
+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 : mergePoints) {
+  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());
     }
   }
 
   allocator.handlePromotionComplete(slot, defaultValue);
 }
 
-LogicalResult SlotPromoter::prepareSlotPromotion() {
-  // First, find the set of operations that will need to be changed for the
-  // promotion to happen. These operations need to resolve some of their uses,
-  // either by rewiring them or simply deleting themselves. If any of them
-  // cannot find a way to resolve their blocking uses, we abort the promotion.
-  if (failed(computeBlockingUses()))
-    return failure();
-
-  // Then, compute blocks in which two or more definitions of the allocated
-  // variable may conflict. These blocks will need a new block argument to
-  // accomodate this.
-  computeMergePoints();
-
-  // The slot can be promoted if the block arguments to be created can
-  // actually be populated with values, which may not be possible depending
-  // on their predecessors.
-  return success(areMergePointsUsable());
-}
-
 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<SlotPromoter> toPromote;
+  SmallVector<MemorySlotPromoter> toPromote;
   for (PromotableAllocationOpInterface allocator : allocators) {
     for (MemorySlot slot : allocator.getPromotableSlots()) {
       if (slot.ptr.use_empty())
         continue;
 
-      SlotPromoter promoter(slot, allocator, builder, dominance);
-      if (succeeded(promoter.prepareSlotPromotion()))
-        toPromote.emplace_back(std::move(promoter));
+      MemorySlotPromotionAnalyzer analyzer(slot, dominance);
+      std::optional<MemorySlotPromotionInfo> info = analyzer.computeInfo();
+      if (info)
+        toPromote.emplace_back(slot, allocator, builder, dominance,
+                               std::move(*info));
     }
   }
 
-  for (SlotPromoter &promoter : toPromote)
+  for (MemorySlotPromoter &promoter : toPromote)
     promoter.promoteSlot();
 
   return success(!toPromote.empty());
 }
 
 namespace {
 
 struct Mem2Reg : impl::Mem2RegBase<Mem2Reg> {