[mlir] Add `requiresReplacedValues` and `visitReplacedValues` to `PromotableOpInterface` #86792

fabianmcg · 2024-03-27T12:00:42Z

Add requiresReplacedValues and visitReplacedValues 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. It is also necessary for enabling #73057

…romotableOpInterface` 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`.

llvmbot · 2024-03-27T12:09:27Z

@llvm/pr-subscribers-mlir-llvm
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Fabian Mora (fabianmcg)

Changes

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.

Full diff: https://github.com/llvm/llvm-project/pull/86792.diff

4 Files Affected:

(modified) mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td (+4-2)
(modified) mlir/include/mlir/Interfaces/MemorySlotInterfaces.td (+21)
(modified) mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp (+12-7)
(modified) mlir/lib/Transforms/Mem2Reg.cpp (+13-1)

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)

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

Moxinilian · 2024-03-27T12:53:16Z

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

+        Checks whether the operation requires visiting the mutated
+        definitions by a store operation.


Could you explain what this means more in the documentation?

Co-authored-by: Théo Degioanni <[email protected]>

Dinistro

I do not see how this is equivalent to the existing implementation.

If I understand correctly, this change will pass all N changed definitions to each visitor-function, instead of running the visitor for each DbgValueOp separately.

Maybe add a test with two dbg value ops on the same store to determine that this works as intended.

fabianmcg · 2024-03-27T13:31:57Z

Maybe add a test with two dbg value ops on the same store to determine that this works as intended.

Are you referring to a test with two dbg declare ops? Because this change should only affect llvm.dbg.declare.
Also, It's worth noting that after the change all existing tests passed.

Dinistro · 2024-03-27T13:44:51Z

Maybe add a test with two dbg value ops on the same store to determine that this works as intended.

Are you referring to a test with two dbg declare ops? Because this change should only affect llvm.dbg.declare. Also, It's worth noting that after the change all existing tests passed.

Sorry for the confusion, yes, I meant two dbg declares.

I'm currently working on extensions for both SROA and Mem2Reg, and I can tell you that the test coverage is unfortunately not sufficient for such cases.

fabianmcg · 2024-03-27T13:51:43Z

@Dinistro because I'm not 100% sure I understand what you are looking for in the test, are you referring to something like:

module {
  llvm.func @basic_store_load(%arg0: i64) -> i64 {
    %0 = llvm.mlir.constant(1 : i32) : i32
    %1 = llvm.alloca %0 x i64 {alignment = 8 : i64} : (i32) -> !llvm.ptr
    llvm.store %arg0, %1 {alignment = 4 : i64} : i64, !llvm.ptr
    llvm.intr.dbg.declare #di_local_variable = %1 : !llvm.ptr
    llvm.intr.dbg.declare #di_local_variable = %1 : !llvm.ptr
    %2 = llvm.load %1 {alignment = 4 : i64} : !llvm.ptr -> i64
    llvm.return %2 : i64
  }
}

WIth result:

module {
  llvm.func @basic_store_load(%arg0: i64) -> i64 {
    %0 = llvm.mlir.constant(1 : i32) : i32
    llvm.intr.dbg.value #di_local_variable = %arg0 : i64
    llvm.intr.dbg.value #di_local_variable = %arg0 : i64
    llvm.return %arg0 : i64
  }
}

Dinistro · 2024-03-27T14:07:29Z

@fabianmcg yes, this is what I'm referring to. Looking at the code indicates to me that this might go wrong, due to running visitMutatedDefs for each declare op but also passing in an ArrayRef. Note that I might also just be confused 🙃

fabianmcg · 2024-03-27T14:30:58Z

@Dinistro I added a test with two llvm.intr.dbg.declare and both versions produce the same result.

Dinistro

Oh, I was indeed confused. Thanks for adding a test for this case.

LGTM, but maybe consider to wait for @Moxinilian to confirm that this is fine for him as well.

Moxinilian · 2024-03-27T16:50:34Z

I'd just like the doc to be a bit clearer about what this is about but otherwise looks good!

…mendMutatedDefs

fabianmcg · 2024-03-27T18:36:17Z

@Moxinilian I improved the docs. I also changed the names of the methods to requiresAmendingMutatedDefs and amendMutatedDefs, as I think the previous names sparked some confusion.

Moxinilian · 2024-03-27T18:48:50Z

I'm sorry but I still don't understand what those methods are about 😅
As a user of the interface, when I see those methods I have no idea what is happening. What does "amending the mutated definitions by a store operation" mean? What is a "mutated definition"?

fabianmcg · 2024-03-28T16:53:12Z

What about:

This method returns whether the promoted operation requires amending the values produced by calls to the PromotableMemOpInterface::getStored method on the slot. Hereafter, these values are called mutated definitions by a store operation.

?

Moxinilian

I think I see what you mean now. Here are some rewording suggestions, I think using simple English helps documentation clarity a lot :)

Moxinilian · 2024-03-28T19:17:58Z

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

+        This method returns whether the promoted operation requires amending the
+        mutated definitions by a store operation.


Suggested change

This method returns whether the promoted operation requires amending the

mutated definitions by a store operation.

This method allows the promoted operation to visit the SSA values used

in place of the memory slot once the promotion process of the memory

slot is complete.

Moxinilian · 2024-03-28T19:19:29Z

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

+        If this method returns true, the operation will be visited using the
+        `amendMutatedDefs` method after the main mutation stage finishes
+        (i.e., after all ops have been processed with `removeBlockingUses`).


Suggested change

If this method returns true, the operation will be visited using the

`amendMutatedDefs` method after the main mutation stage finishes

(i.e., after all ops have been processed with `removeBlockingUses`).

If this method returns true, the `amendMutatedDefs` method on this

operation will be called after the main mutation stage finishes

(i.e., after all ops have been processed with `removeBlockingUses`).

Moxinilian · 2024-03-28T19:26:07Z

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

+        Transforms the IR to amend all the mutated definitions to the slot by a
+        store operation.


Suggested change

Transforms the IR to amend all the mutated definitions to the slot by a

store operation.

Transforms IR using the SSA values that replaced the memory slot.

Moxinilian · 2024-03-28T19:30:31Z

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

+      (ins "::llvm::ArrayRef<std::pair<::mlir::Operation*, ::mlir::Value>>":$mutatedDefs,
+           "::mlir::RewriterBase &":$rewriter), [{}], [{ return; }]


I feel like to be absolutely complete this method should also provide the memory slot, as one may want to know which memory slot specifically those definitions correspond to. But at the same time I am not really sure that this would be useful in practice, so change it only if you feel like it.

Moxinilian · 2024-03-28T19:31:16Z

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

+        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), [{}],


In a perfect world I'd prefer those methods to be called something like visitReplacingValues, because "amending mutated definitions" feels very vague in context.

Moxinilian · 2024-03-28T19:32:05Z

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

+        mutations must happen through the rewriter. During the transformation,
+        *no operation should be deleted*.
+      }],
+      "void", "amendMutatedDefs",


Name suggestion: visitReplacedValues

Moxinilian · 2024-03-28T19:32:28Z

mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

+        *no operation should be deleted*.
+      }],
+      "void", "amendMutatedDefs",
+      (ins "::llvm::ArrayRef<std::pair<::mlir::Operation*, ::mlir::Value>>":$mutatedDefs,


Name suggestion: replacingValues

Moxinilian · 2024-03-28T19:34:53Z

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;


I really dislike the "mutatedValues" name because it is unclear to me in what way those values are mutated. This name suggests you put something in those values when what happens conceptually is the opposite.

Moxinilian · 2024-03-28T19:36:52Z

mlir/lib/Transforms/Mem2Reg.cpp

@@ -552,6 +554,8 @@ void MemorySlotPromoter::removeBlockingUses() {
  dominanceSort(usersToRemoveUses, *slot.ptr.getParentBlock()->getParent());

  llvm::SmallVector<Operation *> toErase;
+  llvm::SmallVector<std::pair<Operation *, Value>> mutatedDefinitions;


Same as above for this name.

fabianmcg · 2024-04-02T13:31:03Z

@Moxinilian I followed all your suggestions but one. I changed requiresAmendingMutatedDefs to requiresReplacedValues, because I felt having visitReplacedValues and visitReplacingValues could create confusion.

Moxinilian

Thank you very much for putting up with my requests ahah

Dinistro · 2024-04-04T07:30:47Z

Can we land this to avoid unnecessary conflicts with potentially upcoming PRs?

fabianmcg · 2024-04-04T15:11:49Z

@Moxinilian, @Dinistro are you both in favor and with high confidence on the change?

Dinistro · 2024-04-04T15:16:54Z

@Moxinilian, @Dinistro are you both in favor and with high confidence on the change?

I think this is fine, feel free to land

Moxinilian · 2024-04-04T16:04:07Z

Same

fabianmcg marked this pull request as ready for review March 27, 2024 12:08

fabianmcg requested a review from gysit March 27, 2024 12:09

llvmbot added mlir:core MLIR Core Infrastructure mlir:llvm mlir labels Mar 27, 2024

fabianmcg requested review from joker-eph and Moxinilian March 27, 2024 12:09

gysit requested review from zyx-billy and Dinistro and removed request for zyx-billy March 27, 2024 12:10

Moxinilian reviewed Mar 27, 2024

View reviewed changes

Update mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

9e85edd

Co-authored-by: Théo Degioanni <[email protected]>

Dinistro reviewed Mar 27, 2024

View reviewed changes

add test and change requiresVisitingMutatedDefs docs

327ee45

Dinistro approved these changes Mar 27, 2024

View reviewed changes

improve docs and change fn names to requiresAmendingMutatedDefs and a…

a03e75d

…mendMutatedDefs

fabianmcg changed the title ~~[mlir] Add requiresVisitingMutatedDefs and visitMutatedDefs to PromotableOpInterface~~ [mlir] Add requiresAmendingMutatedDefs and amendMutatedDefs to PromotableOpInterface Mar 27, 2024

Moxinilian reviewed Mar 28, 2024

View reviewed changes

address reviewer comments

975e1ca

Moxinilian approved these changes Apr 2, 2024

View reviewed changes

fabianmcg changed the title ~~[mlir] Add requiresAmendingMutatedDefs and amendMutatedDefs to PromotableOpInterface~~ [mlir] Add requiresReplacedValues and visitReplacedValues to PromotableOpInterface Apr 2, 2024

fabianmcg mentioned this pull request Apr 4, 2024

[mlir][Ptr] Init the Ptr dialect with the !ptr.ptr type. #86860

Merged

fabianmcg merged commit 220cdf9 into llvm:main Apr 4, 2024

fabianmcg deleted the pr-memory-slot branch August 5, 2024 12:22

		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.

-        This method returns whether the promoted operation requires amending the
-        mutated definitions by a store operation.
+        This method allows the promoted operation to visit the SSA values used
+        in place of the memory slot once the promotion process of the memory
+        slot is complete.

		Transforms the IR to amend all the mutated definitions to the slot by a
		store operation.

	Transforms the IR to amend all the mutated definitions to the slot by a
	store operation.
	Transforms IR using the SSA values that replaced the memory slot.

		(ins "::llvm::ArrayRef<std::pair<::mlir::Operation*, ::mlir::Value>>":$mutatedDefs,
		"::mlir::RewriterBase &":$rewriter), [{}], [{ return; }]

[mlir] Add requiresReplacedValues and visitReplacedValues to PromotableOpInterface #86792

[mlir] Add requiresReplacedValues and visitReplacedValues to PromotableOpInterface #86792

Uh oh!

Conversation

fabianmcg commented Mar 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Mar 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Dinistro left a comment

Choose a reason for hiding this comment

Uh oh!

fabianmcg commented Mar 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Dinistro commented Mar 27, 2024

Uh oh!

fabianmcg commented Mar 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Dinistro commented Mar 27, 2024

Uh oh!

fabianmcg commented Mar 27, 2024

Uh oh!

Dinistro left a comment

Choose a reason for hiding this comment

Uh oh!

Moxinilian commented Mar 27, 2024

Uh oh!

fabianmcg commented Mar 27, 2024

Uh oh!

Moxinilian commented Mar 27, 2024

Uh oh!

fabianmcg commented Mar 28, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Moxinilian left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fabianmcg commented Apr 2, 2024

Uh oh!

Moxinilian left a comment

Choose a reason for hiding this comment

Uh oh!

Dinistro commented Apr 4, 2024

Uh oh!

fabianmcg commented Apr 4, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Dinistro commented Apr 4, 2024

Uh oh!

Moxinilian commented Apr 4, 2024

Uh oh!

Uh oh!

[mlir] Add `requiresReplacedValues` and `visitReplacedValues` to `PromotableOpInterface` #86792

[mlir] Add `requiresReplacedValues` and `visitReplacedValues` to `PromotableOpInterface` #86792

fabianmcg commented Mar 27, 2024 •

edited

Loading

llvmbot commented Mar 27, 2024 •

edited

Loading

fabianmcg commented Mar 27, 2024 •

edited

Loading

fabianmcg commented Mar 27, 2024 •

edited

Loading

fabianmcg commented Mar 28, 2024 •

edited

Loading

fabianmcg commented Apr 4, 2024 •

edited

Loading