[SROA][Mem2Reg] Use efficient droppable use API (after D83976)

Reviewed By: efriedma Differential Revision: https://reviews.llvm.org/D84804
llvm · Jul 28, 2020 · 450dc09 · 450dc09
1 parent 417d3d4
commit 450dc09
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 41 deletions.
diff --git a/llvm/include/llvm/IR/Value.h b/llvm/include/llvm/IR/Value.h
@@ -470,8 +470,11 @@ class Value {
   void dropDroppableUses(llvm::function_ref<bool(const Use *)> ShouldDrop =
                              [](const Use *) { return true; });
 
-  /// Remove every use of \p User that can safely be removed.
-  void dropDroppableUsesByUser(const User &Usr);
+  /// Remove every use of this value in \p User that can safely be removed.
+  void dropDroppableUsesIn(User &Usr);
+
+  /// Remove the droppable use \p U.
+  void dropDroppableUse(Use &U);
 
   /// Check if this value is used in the specified basic block.
   bool isUsedInBasicBlock(const BasicBlock *BB) const;

diff --git a/llvm/lib/IR/Value.cpp b/llvm/lib/IR/Value.cpp
@@ -175,25 +175,34 @@ void Value::dropDroppableUses(
   for (Use &U : uses())
     if (U.getUser()->isDroppable() && ShouldDrop(&U))
       ToBeEdited.push_back(&U);
-  for (Use *U : ToBeEdited) {
-    U->removeFromList();
-    if (auto *Assume = dyn_cast<IntrinsicInst>(U->getUser())) {
-      assert(Assume->getIntrinsicID() == Intrinsic::assume);
-      unsigned OpNo = U->getOperandNo();
-      if (OpNo == 0)
-        Assume->setOperand(0, ConstantInt::getTrue(Assume->getContext()));
-      else {
-        Assume->setOperand(OpNo, UndefValue::get(U->get()->getType()));
-        CallInst::BundleOpInfo &BOI = Assume->getBundleOpInfoForOperand(OpNo);
-        BOI.Tag = getContext().pImpl->getOrInsertBundleTag("ignore");
-      }
-    } else
-      llvm_unreachable("unkown droppable use");
+  for (Use *U : ToBeEdited)
+    dropDroppableUse(*U);
+}
+
+void Value::dropDroppableUsesIn(User &Usr) {
+  assert(Usr.isDroppable() && "Expected a droppable user!");
+  for (Use &UsrOp : Usr.operands()) {
+    if (UsrOp.get() == this)
+      dropDroppableUse(UsrOp);
   }
 }
 
-void Value::dropDroppableUsesByUser(const User &Usr) {
-  dropDroppableUses([&](const Use *U) { return U->getUser() == &Usr; });
+void Value::dropDroppableUse(Use &U) {
+  U.removeFromList();
+  if (auto *Assume = dyn_cast<IntrinsicInst>(U.getUser())) {
+    assert(Assume->getIntrinsicID() == Intrinsic::assume);
+    unsigned OpNo = U.getOperandNo();
+    if (OpNo == 0)
+      U.set(ConstantInt::getTrue(Assume->getContext()));
+    else {
+      U.set(UndefValue::get(U.get()->getType()));
+      CallInst::BundleOpInfo &BOI = Assume->getBundleOpInfoForOperand(OpNo);
+      BOI.Tag = getContext().pImpl->getOrInsertBundleTag("ignore");
+    }
+    return;
+  }
+
+  llvm_unreachable("unkown droppable use");
 }
 
 bool Value::isUsedInBasicBlock(const BasicBlock *BB) const {

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -3087,7 +3087,7 @@ class llvm::sroa::AllocaSliceRewriter
     if (II.isDroppable()) {
       assert(II.getIntrinsicID() == Intrinsic::assume && "Expected assume");
       // TODO For now we forget assumed information, this can be improved.
-      OldPtr->dropDroppableUsesByUser(II);
+      OldPtr->dropDroppableUsesIn(II);
       return true;
     }
 

diff --git a/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp b/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
@@ -311,45 +311,31 @@ static void removeIntrinsicUsers(AllocaInst *AI) {
   // Knowing that this alloca is promotable, we know that it's safe to kill all
   // instructions except for load and store.
 
-  // Helper to drop the uses of \p I in \p UserI.
-  auto DropUsesIn = [](Instruction *UserI, Instruction *I,
-                       Instruction::user_iterator &UI,
-                       const Instruction::user_iterator &UE) {
-    // TODO For now we forget assumed information, this can be improved.
-    assert(isa<IntrinsicInst>(UserI) &&
-           cast<IntrinsicInst>(UserI)->getIntrinsicID() == Intrinsic::assume &&
-           "Expected assume");
-
-    // Skip ahead if User has multiple uses of I.
-    while (UI != UE && *UI == UserI)
-      ++UI;
-
-    I->dropDroppableUsesByUser(*UserI);
-  };
-
-  for (auto UI = AI->user_begin(), UE = AI->user_end(); UI != UE;) {
-    Instruction *I = cast<Instruction>(*UI);
+  for (auto UI = AI->use_begin(), UE = AI->use_end(); UI != UE;) {
+    Instruction *I = cast<Instruction>(UI->getUser());
+    Use &U = *UI;
     ++UI;
     if (isa<LoadInst>(I) || isa<StoreInst>(I))
       continue;
 
     // Drop the use of AI in droppable instructions.
     if (I->isDroppable()) {
-      DropUsesIn(I, AI, UI, UE);
+      I->dropDroppableUse(U);
       continue;
     }
 
     if (!I->getType()->isVoidTy()) {
       // The only users of this bitcast/GEP instruction are lifetime intrinsics.
       // Follow the use/def chain to erase them now instead of leaving it for
       // dead code elimination later.
-      for (auto UUI = I->user_begin(), UUE = I->user_end(); UUI != UUE;) {
-        Instruction *Inst = cast<Instruction>(*UUI);
+      for (auto UUI = I->use_begin(), UUE = I->use_end(); UUI != UUE;) {
+        Instruction *Inst = cast<Instruction>(UUI->getUser());
+        Use &UU = *UUI;
         ++UUI;
 
         // Drop the use of I in droppable instructions.
         if (Inst->isDroppable()) {
-          DropUsesIn(Inst, I, UUI, UUE);
+          Inst->dropDroppableUse(UU);
           continue;
         }
         Inst->eraseFromParent();