[Mlir-commits] [mlir] [mlir] Add `requiresVisitingMutatedDefs` and `visitMutatedDefs` to `PromotableOpInterface` (PR #86792)
Fabian Mora
llvmlistbot at llvm.org
Wed Mar 27 11:34:38 PDT 2024
https://github.com/fabianmcg updated https://github.com/llvm/llvm-project/pull/86792
>From e3b7f23fa27f6c28f017b62639675be1926a834f Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Wed, 27 Mar 2024 11:55:39 +0000
Subject: [PATCH 1/4] [mlir] Add `requiresVisitingMutatedDefs` and
`visitMutatedDefs` to `PromotableOpInterface`
Add `requiresVisitingMutatedDefs` and `visitMutatedDefs` methods to
`PromotableOpInterface`. These methods allow `PromotableOpInterface` ops to
transforms definitions mutated by a `store`.
This change is necessary to correctly handle the promotion of
`LLVM_DbgDeclareOp`.
---
.../mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td | 6 ++++--
.../mlir/Interfaces/MemorySlotInterfaces.td | 21 +++++++++++++++++++
mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp | 19 ++++++++++-------
mlir/lib/Transforms/Mem2Reg.cpp | 14 ++++++++++++-
4 files changed, 50 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index f4bac9376f2ea0..d3857fbe06e355 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -562,8 +562,10 @@ class LLVM_DbgIntrOp<string name, string argName, list<Trait> traits = []>
}];
}
-def LLVM_DbgDeclareOp : LLVM_DbgIntrOp<"dbg.declare", "addr",
- [DeclareOpInterfaceMethods<PromotableOpInterface>]> {
+def LLVM_DbgDeclareOp : LLVM_DbgIntrOp<"dbg.declare", "addr", [
+ DeclareOpInterfaceMethods<PromotableOpInterface, [
+ "requiresVisitingMutatedDefs", "visitMutatedDefs"
+ ]>]> {
let summary = "Describes how the address relates to a source language variable.";
let arguments = (ins
LLVM_AnyPointer:$addr,
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index e10e2d4e104c3e..af120eb059aebf 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
@@ -229,6 +229,27 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
(ins "const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
"::mlir::RewriterBase &":$rewriter)
>,
+ InterfaceMethod<[{
+ Checks whether the operation requires visiting the mutated
+ definitions by a store operation.
+ }], "bool", "requiresVisitingMutatedDefs", (ins), [{}],
+ [{ return false; }]
+ >,
+ InterfaceMethod<[{
+ Visits all the mutated definitions by a store operation.
+
+ This method will only be called after all blocking issues haven been
+ scheduled for removal and if `requiresVisitingMutatedDefs` returned
+ true.
+
+ The rewriter is located after the promotable operation on call. All IR
+ mutations must happen through the rewriter. During the transformation,
+ *no operation should be deleted*.
+ }],
+ "void", "visitMutatedDefs",
+ (ins "::llvm::ArrayRef<std::pair<::mlir::Operation*, ::mlir::Value>>":$mutatedDefs,
+ "::mlir::RewriterBase &":$rewriter), [{}], [{ return; }]
+ >,
];
}
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index f171bf7cc4bec3..b5c5f4454dede9 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -166,13 +166,6 @@ bool LLVM::StoreOp::canUsesBeRemoved(
DeletionKind LLVM::StoreOp::removeBlockingUses(
const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
RewriterBase &rewriter, Value reachingDefinition) {
- // `canUsesBeRemoved` checked this blocking use must be the stored slot
- // pointer.
- for (Operation *user : slot.ptr.getUsers())
- if (auto declareOp = dyn_cast<LLVM::DbgDeclareOp>(user))
- rewriter.create<LLVM::DbgValueOp>(declareOp->getLoc(), getValue(),
- declareOp.getVarInfo(),
- declareOp.getLocationExpr());
return DeletionKind::Delete;
}
@@ -405,6 +398,18 @@ DeletionKind LLVM::DbgValueOp::removeBlockingUses(
return DeletionKind::Keep;
}
+bool LLVM::DbgDeclareOp::requiresVisitingMutatedDefs() { return true; }
+
+void LLVM::DbgDeclareOp::visitMutatedDefs(
+ ArrayRef<std::pair<Operation *, Value>> definitions,
+ RewriterBase &rewriter) {
+ for (auto [op, value] : definitions) {
+ rewriter.setInsertionPointAfter(op);
+ rewriter.create<LLVM::DbgValueOp>(getLoc(), value, getVarInfo(),
+ getLocationExpr());
+ }
+}
+
//===----------------------------------------------------------------------===//
// Interfaces for GEPOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 80e3b790163297..3aca5eb19ff5af 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -202,6 +202,7 @@ class MemorySlotPromoter {
/// Contains the reaching definition at this operation. Reaching definitions
/// are only computed for promotable memory operations with blocking uses.
DenseMap<PromotableMemOpInterface, Value> reachingDefs;
+ DenseMap<PromotableMemOpInterface, Value> mutatedValues;
DominanceInfo &dominance;
MemorySlotPromotionInfo info;
const Mem2RegStatistics &statistics;
@@ -438,6 +439,7 @@ Value MemorySlotPromoter::computeReachingDefInBlock(Block *block,
assert(stored && "a memory operation storing to a slot must provide a "
"new definition of the slot");
reachingDef = stored;
+ mutatedValues[memOp] = stored;
}
}
}
@@ -552,6 +554,8 @@ void MemorySlotPromoter::removeBlockingUses() {
dominanceSort(usersToRemoveUses, *slot.ptr.getParentBlock()->getParent());
llvm::SmallVector<Operation *> toErase;
+ llvm::SmallVector<std::pair<Operation *, Value>> mutatedDefinitions;
+ llvm::SmallVector<PromotableOpInterface> visitMutatedDefinitions;
for (Operation *toPromote : llvm::reverse(usersToRemoveUses)) {
if (auto toPromoteMemOp = dyn_cast<PromotableMemOpInterface>(toPromote)) {
Value reachingDef = reachingDefs.lookup(toPromoteMemOp);
@@ -565,7 +569,9 @@ void MemorySlotPromoter::removeBlockingUses() {
slot, info.userToBlockingUses[toPromote], rewriter,
reachingDef) == DeletionKind::Delete)
toErase.push_back(toPromote);
-
+ if (toPromoteMemOp.storesTo(slot))
+ if (Value mutatedValue = mutatedValues[toPromoteMemOp])
+ mutatedDefinitions.push_back({toPromoteMemOp, mutatedValue});
continue;
}
@@ -574,6 +580,12 @@ void MemorySlotPromoter::removeBlockingUses() {
if (toPromoteBasic.removeBlockingUses(info.userToBlockingUses[toPromote],
rewriter) == DeletionKind::Delete)
toErase.push_back(toPromote);
+ if (toPromoteBasic.requiresVisitingMutatedDefs())
+ visitMutatedDefinitions.push_back(toPromoteBasic);
+ }
+ for (PromotableOpInterface op : visitMutatedDefinitions) {
+ rewriter.setInsertionPointAfter(op);
+ op.visitMutatedDefs(mutatedDefinitions, rewriter);
}
for (Operation *toEraseOp : toErase)
>From 9e85edd6c7ad0d9b5abf3031a25f3d067d9183fc Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Wed, 27 Mar 2024 09:10:14 -0400
Subject: [PATCH 2/4] Update
mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Co-authored-by: Théo Degioanni <30992420+Moxinilian at users.noreply.github.com>
---
mlir/include/mlir/Interfaces/MemorySlotInterfaces.td | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index af120eb059aebf..abc333f3b10f4c 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
@@ -238,7 +238,7 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
InterfaceMethod<[{
Visits all the mutated definitions by a store operation.
- This method will only be called after all blocking issues haven been
+ This method will only be called after all blocking uses have been
scheduled for removal and if `requiresVisitingMutatedDefs` returned
true.
>From 327ee452a682aef84177709e925912c61f6abb3e Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Wed, 27 Mar 2024 14:21:19 +0000
Subject: [PATCH 3/4] add test and change requiresVisitingMutatedDefs docs
---
.../mlir/Interfaces/MemorySlotInterfaces.td | 3 +++
mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir | 21 +++++++++++++++++++
2 files changed, 24 insertions(+)
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index abc333f3b10f4c..308949b9fe6df9 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
@@ -232,6 +232,9 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
InterfaceMethod<[{
Checks whether the operation requires visiting the mutated
definitions by a store operation.
+ If this method returns true, the operation will be visited using the
+ `visitMutatedDefs` method after the main mutation stage finishes
+ (i.e., after all ops have been processed with `removeBlockingUses`).
}], "bool", "requiresVisitingMutatedDefs", (ins), [{}],
[{ return false; }]
>,
diff --git a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
index f7ddb4a7abe5ad..b7cbd787f06e44 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
@@ -29,6 +29,27 @@ llvm.func @basic_store_load(%arg0: i64) -> i64 {
llvm.return %2 : i64
}
+// CHECK-LABEL: llvm.func @multiple_store_load
+llvm.func @multiple_store_load(%arg0: i64) -> i64 {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ // CHECK-NOT: = llvm.alloca
+ %1 = llvm.alloca %0 x i64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
+ // CHECK-NOT: llvm.intr.dbg.declare
+ llvm.intr.dbg.declare #di_local_variable = %1 : !llvm.ptr
+ // CHECK-NOT: llvm.store
+ llvm.store %arg0, %1 {alignment = 4 : i64} : i64, !llvm.ptr
+ // CHECK-NOT: llvm.intr.dbg.declare
+ llvm.intr.dbg.declare #di_local_variable = %1 : !llvm.ptr
+ // CHECK: llvm.intr.dbg.value #[[$VAR]] = %[[LOADED:.*]] : i64
+ // CHECK: llvm.intr.dbg.value #[[$VAR]] = %[[LOADED]] : i64
+ // CHECK-NOT: llvm.intr.dbg.value
+ // CHECK-NOT: llvm.intr.dbg.declare
+ // CHECK-NOT: llvm.store
+ %2 = llvm.load %1 {alignment = 4 : i64} : !llvm.ptr -> i64
+ // CHECK: llvm.return %[[LOADED]] : i64
+ llvm.return %2 : i64
+}
+
// CHECK-LABEL: llvm.func @block_argument_value
// CHECK-SAME: (%[[ARG0:.*]]: i64, {{.*}})
llvm.func @block_argument_value(%arg0: i64, %arg1: i1) -> i64 {
>From a03e75d1bab66d77037179ba7251bece15d4cc87 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Wed, 27 Mar 2024 18:30:17 +0000
Subject: [PATCH 4/4] improve docs and change fn names to
requiresAmendingMutatedDefs and amendMutatedDefs
---
.../mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td | 2 +-
.../mlir/Interfaces/MemorySlotInterfaces.td | 20 ++++++++++++-------
mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp | 4 ++--
mlir/lib/Transforms/Mem2Reg.cpp | 10 +++++-----
4 files changed, 21 insertions(+), 15 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index d3857fbe06e355..e3b53f8ffe835c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -564,7 +564,7 @@ class LLVM_DbgIntrOp<string name, string argName, list<Trait> traits = []>
def LLVM_DbgDeclareOp : LLVM_DbgIntrOp<"dbg.declare", "addr", [
DeclareOpInterfaceMethods<PromotableOpInterface, [
- "requiresVisitingMutatedDefs", "visitMutatedDefs"
+ "requiresAmendingMutatedDefs", "amendMutatedDefs"
]>]> {
let summary = "Describes how the address relates to a source language variable.";
let arguments = (ins
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index 308949b9fe6df9..6aeccd49e78ba9 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
@@ -230,26 +230,32 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
"::mlir::RewriterBase &":$rewriter)
>,
InterfaceMethod<[{
- Checks whether the operation requires visiting the mutated
- definitions by a store operation.
+ This method returns whether the promoted operation requires amending the
+ mutated definitions by a store operation.
+
If this method returns true, the operation will be visited using the
- `visitMutatedDefs` method after the main mutation stage finishes
+ `amendMutatedDefs` method after the main mutation stage finishes
(i.e., after all ops have been processed with `removeBlockingUses`).
- }], "bool", "requiresVisitingMutatedDefs", (ins), [{}],
+
+ Operations should only request amending the mutated definitions if the
+ intended transformation applies to all mutated values. Furthermore,
+ mutated values must not be deleted.
+ }], "bool", "requiresAmendingMutatedDefs", (ins), [{}],
[{ return false; }]
>,
InterfaceMethod<[{
- Visits all the mutated definitions by a store operation.
+ Transforms the IR to amend all the mutated definitions to the slot by a
+ store operation.
This method will only be called after all blocking uses have been
- scheduled for removal and if `requiresVisitingMutatedDefs` returned
+ scheduled for removal and if `requiresAmendingMutatedDefs` returned
true.
The rewriter is located after the promotable operation on call. All IR
mutations must happen through the rewriter. During the transformation,
*no operation should be deleted*.
}],
- "void", "visitMutatedDefs",
+ "void", "amendMutatedDefs",
(ins "::llvm::ArrayRef<std::pair<::mlir::Operation*, ::mlir::Value>>":$mutatedDefs,
"::mlir::RewriterBase &":$rewriter), [{}], [{ return; }]
>,
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index b5c5f4454dede9..4783d4c11755bc 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -398,9 +398,9 @@ DeletionKind LLVM::DbgValueOp::removeBlockingUses(
return DeletionKind::Keep;
}
-bool LLVM::DbgDeclareOp::requiresVisitingMutatedDefs() { return true; }
+bool LLVM::DbgDeclareOp::requiresAmendingMutatedDefs() { return true; }
-void LLVM::DbgDeclareOp::visitMutatedDefs(
+void LLVM::DbgDeclareOp::amendMutatedDefs(
ArrayRef<std::pair<Operation *, Value>> definitions,
RewriterBase &rewriter) {
for (auto [op, value] : definitions) {
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 3aca5eb19ff5af..609a4e593eb302 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -555,7 +555,7 @@ void MemorySlotPromoter::removeBlockingUses() {
llvm::SmallVector<Operation *> toErase;
llvm::SmallVector<std::pair<Operation *, Value>> mutatedDefinitions;
- llvm::SmallVector<PromotableOpInterface> visitMutatedDefinitions;
+ llvm::SmallVector<PromotableOpInterface> toAmend;
for (Operation *toPromote : llvm::reverse(usersToRemoveUses)) {
if (auto toPromoteMemOp = dyn_cast<PromotableMemOpInterface>(toPromote)) {
Value reachingDef = reachingDefs.lookup(toPromoteMemOp);
@@ -580,12 +580,12 @@ void MemorySlotPromoter::removeBlockingUses() {
if (toPromoteBasic.removeBlockingUses(info.userToBlockingUses[toPromote],
rewriter) == DeletionKind::Delete)
toErase.push_back(toPromote);
- if (toPromoteBasic.requiresVisitingMutatedDefs())
- visitMutatedDefinitions.push_back(toPromoteBasic);
+ if (toPromoteBasic.requiresAmendingMutatedDefs())
+ toAmend.push_back(toPromoteBasic);
}
- for (PromotableOpInterface op : visitMutatedDefinitions) {
+ for (PromotableOpInterface op : toAmend) {
rewriter.setInsertionPointAfter(op);
- op.visitMutatedDefs(mutatedDefinitions, rewriter);
+ op.amendMutatedDefs(mutatedDefinitions, rewriter);
}
for (Operation *toEraseOp : toErase)
More information about the Mlir-commits
mailing list