Skip to content

Commit 92cc30a

Browse files
Théo Degioannigysit
authored andcommitted
Reland "[mlir][mem2reg] Expose algorithm internals."
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. This commit reverts f333977 and relands 91cff8a. The original commit was reverted accidentally due to a misinterpretation of a bazel build bot failure. Reviewed By: gysit Differential Revision: https://reviews.llvm.org/D149825
1 parent feb69da commit 92cc30a

File tree

15 files changed

+242
-202
lines changed

15 files changed

+242
-202
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ include "mlir/Dialect/LLVMIR/LLVMAttrDefs.td"
66
include "mlir/Dialect/LLVMIR/LLVMEnums.td"
77
include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
88
include "mlir/Interfaces/InferTypeOpInterface.td"
9-
include "mlir/Interfaces/Mem2RegInterfaces.td"
9+
include "mlir/Interfaces/MemorySlotInterfaces.td"
1010

1111
// Operations that correspond to LLVM intrinsics. With MLIR operation set being
1212
// extendable, there is no reason to introduce a hard boundary between "core"

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ include "mlir/IR/SymbolInterfaces.td"
2222
include "mlir/Interfaces/CallInterfaces.td"
2323
include "mlir/Interfaces/ControlFlowInterfaces.td"
2424
include "mlir/Interfaces/InferTypeOpInterface.td"
25-
include "mlir/Interfaces/Mem2RegInterfaces.td"
25+
include "mlir/Interfaces/MemorySlotInterfaces.td"
2626
include "mlir/Interfaces/SideEffectInterfaces.td"
2727

2828
class LLVM_Builder<string builder> {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
#include "mlir/Interfaces/ControlFlowInterfaces.h"
1717
#include "mlir/Interfaces/CopyOpInterface.h"
1818
#include "mlir/Interfaces/InferTypeOpInterface.h"
19-
#include "mlir/Interfaces/Mem2RegInterfaces.h"
19+
#include "mlir/Interfaces/MemorySlotInterfaces.h"
2020
#include "mlir/Interfaces/ShapedOpInterfaces.h"
2121
#include "mlir/Interfaces/SideEffectInterfaces.h"
2222
#include "mlir/Interfaces/ViewLikeInterface.h"

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ include "mlir/Interfaces/CastInterfaces.td"
1515
include "mlir/Interfaces/ControlFlowInterfaces.td"
1616
include "mlir/Interfaces/CopyOpInterface.td"
1717
include "mlir/Interfaces/InferTypeOpInterface.td"
18-
include "mlir/Interfaces/Mem2RegInterfaces.td"
18+
include "mlir/Interfaces/MemorySlotInterfaces.td"
1919
include "mlir/Interfaces/ShapedOpInterfaces.td"
2020
include "mlir/Interfaces/SideEffectInterfaces.td"
2121
include "mlir/Interfaces/ViewLikeInterface.td"

mlir/include/mlir/Interfaces/CMakeLists.txt

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ add_mlir_interface(DestinationStyleOpInterface)
77
add_mlir_interface(InferIntRangeInterface)
88
add_mlir_interface(InferTypeOpInterface)
99
add_mlir_interface(LoopLikeInterface)
10-
add_mlir_interface(Mem2RegInterfaces)
1110
add_mlir_interface(ParallelCombiningOpInterface)
1211
add_mlir_interface(RuntimeVerifiableOpInterface)
1312
add_mlir_interface(ShapedOpInterfaces)
@@ -17,6 +16,12 @@ add_mlir_interface(ValueBoundsOpInterface)
1716
add_mlir_interface(VectorInterfaces)
1817
add_mlir_interface(ViewLikeInterface)
1918

19+
set(LLVM_TARGET_DEFINITIONS MemorySlotInterfaces.td)
20+
mlir_tablegen(MemorySlotOpInterfaces.h.inc -gen-op-interface-decls)
21+
mlir_tablegen(MemorySlotOpInterfaces.cpp.inc -gen-op-interface-defs)
22+
add_public_tablegen_target(MLIRMemorySlotInterfacesIncGen)
23+
add_dependencies(mlir-generic-headers MLIRMemorySlotInterfacesIncGen)
24+
2025
set(LLVM_TARGET_DEFINITIONS DataLayoutInterfaces.td)
2126
mlir_tablegen(DataLayoutAttrInterface.h.inc -gen-attr-interface-decls)
2227
mlir_tablegen(DataLayoutAttrInterface.cpp.inc -gen-attr-interface-defs)

mlir/include/mlir/Interfaces/Mem2RegInterfaces.h renamed to mlir/include/mlir/Interfaces/MemorySlotInterfaces.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
//
77
//===----------------------------------------------------------------------===//
88

9-
#ifndef MLIR_INTERFACES_MEM2REGINTERFACES_H
10-
#define MLIR_INTERFACES_MEM2REGINTERFACES_H
9+
#ifndef MLIR_INTERFACES_MEMORYSLOTINTERFACES_H
10+
#define MLIR_INTERFACES_MEMORYSLOTINTERFACES_H
1111

1212
#include "mlir/IR/Dominance.h"
1313
#include "mlir/IR/OpDefinition.h"
@@ -34,6 +34,6 @@ enum class DeletionKind {
3434

3535
} // namespace mlir
3636

37-
#include "mlir/Interfaces/Mem2RegInterfaces.h.inc"
37+
#include "mlir/Interfaces/MemorySlotOpInterfaces.h.inc"
3838

39-
#endif // MLIR_INTERFACES_MEM2REGINTERFACES_H
39+
#endif // MLIR_INTERFACES_MEMORYSLOTINTERFACES_H

mlir/include/mlir/Interfaces/Mem2RegInterfaces.td renamed to mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
1-
//===-- Mem2RegInterfaces.td - Mem2Reg interfaces ----------*- tablegen -*-===//
1+
//===-- MemorySlotInterfaces.td - MemorySlot interfaces ----*- tablegen -*-===//
22
//
33
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
44
// See https://llvm.org/LICENSE.txt for license information.
55
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
66
//
77
//===----------------------------------------------------------------------===//
88

9-
#ifndef MLIR_INTERFACES_MEM2REGINTERFACES
10-
#define MLIR_INTERFACES_MEM2REGINTERFACES
9+
#ifndef MLIR_INTERFACES_MEMORYSLOTINTERFACES
10+
#define MLIR_INTERFACES_MEMORYSLOTINTERFACES
1111

1212
include "mlir/IR/OpBase.td"
1313

@@ -76,6 +76,9 @@ def PromotableMemOpInterface : OpInterface<"PromotableMemOpInterface"> {
7676
to memory slots. Loads and stores must be of whole values of the same
7777
type as the slot itself.
7878

79+
For a memory operation on a slot to be valid, it must operate on the slot
80+
pointer *only as a pointer to an element of the type of the slot*.
81+
7982
If the same operation does both loads and stores on the same slot, the
8083
load must semantically happen first.
8184
}];
@@ -152,21 +155,21 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
152155
let methods = [
153156
InterfaceMethod<[{
154157
Checks that this operation can be promoted to no longer use the provided
155-
blocking uses, in the context of promoting `slot`.
158+
blocking uses, in order to allow optimization.
156159

157160
If the removal procedure of the use will require that other uses get
158161
removed, that dependency should be added to the `newBlockingUses`
159162
argument. Dependent uses must only be uses of results of this operation.
160163
}], "bool", "canUsesBeRemoved",
161-
(ins "const ::mlir::MemorySlot &":$slot,
162-
"const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &":$blockingUses,
164+
(ins "const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &":$blockingUses,
163165
"::llvm::SmallVectorImpl<::mlir::OpOperand *> &":$newBlockingUses)
164166
>,
165167
InterfaceMethod<[{
166168
Transforms IR to ensure that the current operation does not use the
167-
provided memory slot anymore. In contrast to `PromotableMemOpInterface`,
168-
operations implementing this interface must not need access to the
169-
reaching definition of the content of the slot.
169+
provided blocking uses anymore. In contrast to
170+
`PromotableMemOpInterface`, operations implementing this interface
171+
must not need access to the reaching definition of the content of the
172+
slot.
170173

171174
During the transformation, *no operation should be deleted*.
172175
The operation can only schedule its own deletion by returning the
@@ -186,11 +189,10 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
186189
}],
187190
"::mlir::DeletionKind",
188191
"removeBlockingUses",
189-
(ins "const ::mlir::MemorySlot &":$slot,
190-
"const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
192+
(ins "const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
191193
"::mlir::OpBuilder &":$builder)
192194
>,
193195
];
194196
}
195197

196-
#endif // MLIR_INTERFACES_MEM2REGINTERFACES
198+
#endif // MLIR_INTERFACES_MEMORYSLOTINTERFACES

mlir/include/mlir/Transforms/Mem2Reg.h

Lines changed: 103 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,112 @@
1111

1212
#include "mlir/IR/Dominance.h"
1313
#include "mlir/IR/OpDefinition.h"
14-
#include "mlir/Interfaces/Mem2RegInterfaces.h"
14+
#include "mlir/Interfaces/MemorySlotInterfaces.h"
1515

1616
namespace mlir {
1717

18+
/// Information computed during promotion analysis used to perform actual
19+
/// promotion.
20+
struct MemorySlotPromotionInfo {
21+
/// Blocks for which at least two definitions of the slot values clash.
22+
SmallPtrSet<Block *, 8> mergePoints;
23+
/// Contains, for each operation, which uses must be eliminated by promotion.
24+
/// This is a DAG structure because if an operation must eliminate some of
25+
/// its uses, it is because the defining ops of the blocking uses requested
26+
/// it. The defining ops therefore must also have blocking uses or be the
27+
/// starting point of the bloccking uses.
28+
DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> userToBlockingUses;
29+
};
30+
31+
/// Computes information for basic slot promotion. This will check that direct
32+
/// slot promotion can be performed, and provide the information to execute the
33+
/// promotion. This does not mutate IR.
34+
class MemorySlotPromotionAnalyzer {
35+
public:
36+
MemorySlotPromotionAnalyzer(MemorySlot slot, DominanceInfo &dominance)
37+
: slot(slot), dominance(dominance) {}
38+
39+
/// Computes the information for slot promotion if promotion is possible,
40+
/// returns nothing otherwise.
41+
std::optional<MemorySlotPromotionInfo> computeInfo();
42+
43+
private:
44+
/// Computes the transitive uses of the slot that block promotion. This finds
45+
/// uses that would block the promotion, checks that the operation has a
46+
/// solution to remove the blocking use, and potentially forwards the analysis
47+
/// if the operation needs further blocking uses resolved to resolve its own
48+
/// uses (typically, removing its users because it will delete itself to
49+
/// resolve its own blocking uses). This will fail if one of the transitive
50+
/// users cannot remove a requested use, and should prevent promotion.
51+
LogicalResult computeBlockingUses(
52+
DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> &userToBlockingUses);
53+
54+
/// Computes in which blocks the value stored in the slot is actually used,
55+
/// meaning blocks leading to a load. This method uses `definingBlocks`, the
56+
/// set of blocks containing a store to the slot (defining the value of the
57+
/// slot).
58+
SmallPtrSet<Block *, 16>
59+
computeSlotLiveIn(SmallPtrSetImpl<Block *> &definingBlocks);
60+
61+
/// Computes the points in which multiple re-definitions of the slot's value
62+
/// (stores) may conflict.
63+
void computeMergePoints(SmallPtrSetImpl<Block *> &mergePoints);
64+
65+
/// Ensures predecessors of merge points can properly provide their current
66+
/// definition of the value stored in the slot to the merge point. This can
67+
/// notably be an issue if the terminator used does not have the ability to
68+
/// forward values through block operands.
69+
bool areMergePointsUsable(SmallPtrSetImpl<Block *> &mergePoints);
70+
71+
MemorySlot slot;
72+
DominanceInfo &dominance;
73+
};
74+
75+
/// The MemorySlotPromoter handles the state of promoting a memory slot. It
76+
/// wraps a slot and its associated allocator. This will perform the mutation of
77+
/// IR.
78+
class MemorySlotPromoter {
79+
public:
80+
MemorySlotPromoter(MemorySlot slot, PromotableAllocationOpInterface allocator,
81+
OpBuilder &builder, DominanceInfo &dominance,
82+
MemorySlotPromotionInfo info);
83+
84+
/// Actually promotes the slot by mutating IR. Promoting a slot does not
85+
/// invalidate the MemorySlotPromotionInfo of other slots.
86+
void promoteSlot();
87+
88+
private:
89+
/// Computes the reaching definition for all the operations that require
90+
/// promotion. `reachingDef` is the value the slot should contain at the
91+
/// beginning of the block. This method returns the reached definition at the
92+
/// end of the block.
93+
Value computeReachingDefInBlock(Block *block, Value reachingDef);
94+
95+
/// Computes the reaching definition for all the operations that require
96+
/// promotion. `reachingDef` corresponds to the initial value the
97+
/// slot will contain before any write, typically a poison value.
98+
void computeReachingDefInRegion(Region *region, Value reachingDef);
99+
100+
/// Removes the blocking uses of the slot, in topological order.
101+
void removeBlockingUses();
102+
103+
/// Lazily-constructed default value representing the content of the slot when
104+
/// no store has been executed. This function may mutate IR.
105+
Value getLazyDefaultValue();
106+
107+
MemorySlot slot;
108+
PromotableAllocationOpInterface allocator;
109+
OpBuilder &builder;
110+
/// Potentially non-initialized default value. Use `getLazyDefaultValue` to
111+
/// initialize it on demand.
112+
Value defaultValue;
113+
/// Contains the reaching definition at this operation. Reaching definitions
114+
/// are only computed for promotable memory operations with blocking uses.
115+
DenseMap<PromotableMemOpInterface, Value> reachingDefs;
116+
DominanceInfo &dominance;
117+
MemorySlotPromotionInfo info;
118+
};
119+
18120
/// Attempts to promote the memory slots of the provided allocators. Succeeds if
19121
/// at least one memory slot was promoted.
20122
LogicalResult

mlir/lib/Dialect/LLVMIR/CMakeLists.txt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ add_mlir_dialect_library(MLIRLLVMDialect
66
IR/LLVMDialect.cpp
77
IR/LLVMInlining.cpp
88
IR/LLVMInterfaces.cpp
9-
IR/LLVMMem2Reg.cpp
9+
IR/LLVMMemorySlot.cpp
1010
IR/LLVMTypes.cpp
1111
IR/LLVMTypeSyntax.cpp
1212

@@ -34,7 +34,7 @@ add_mlir_dialect_library(MLIRLLVMDialect
3434
MLIRDataLayoutInterfaces
3535
MLIRInferTypeOpInterface
3636
MLIRIR
37-
MLIRMem2RegInterfaces
37+
MLIRMemorySlotInterfaces
3838
MLIRSideEffectInterfaces
3939
MLIRSupport
4040
)

0 commit comments

Comments
 (0)