[Attributor][Fix] Resolve various liveness issues

Summary:
This contains various fixes:
  - Explicitly determine and return the next noreturn instruction.
  - If an invoke calls a noreturn function which is not nounwind we
    keep the unwind destination live. This also means we require an
    invoke. Though we can still add the unreachable to the normal
    destination block.
  - Check if the return instructions are dead after we look for calls
    to avoid triggering an optimistic fixpoint in the presence of
    assumed liveness information.
  - Make the interface work with "const" pointers.
  - Some simplifications

While additional tests are included, full coverage is achieved only with
D59978.

Reviewers: sstefan1, uenoku

Subscribers: hiraditya, bollu, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D65701

llvm-svn: 367791

1 files changed, 94 insertions, 72 deletions

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 87bdc0d0aa5..1321638d94c 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -753,11 +753,6 @@ ChangeStatus AAReturnedValuesImpl::updateImpl(Attributor &A) {
     SmallPtrSet<ReturnInst *, 2> &ReturnInsts = It.second;
     Value *RV = It.first;
 
-    // Ignore dead ReturnValues.
-    if (LivenessAA &&
-        !LivenessAA->isLiveInstSet(ReturnInsts.begin(), ReturnInsts.end()))
-      continue;
-
     LLVM_DEBUG(dbgs() << "[AAReturnedValues] Potentially returned value " << *RV
                       << "\n");
 
@@ -771,6 +766,14 @@ ChangeStatus AAReturnedValuesImpl::updateImpl(Attributor &A) {
     // sites will be removed and we will fix the information for this state.
     HasCallSite = true;
 
+    // Ignore dead ReturnValues.
+    if (LivenessAA &&
+        !LivenessAA->isLiveInstSet(ReturnInsts.begin(), ReturnInsts.end())) {
+      LLVM_DEBUG(dbgs() << "[AAReturnedValues] all returns are assumed dead, "
+                           "skip it for now\n");
+      continue;
+    }
+
     // Try to find a assumed unique return value for the called function.
     auto *RetCSAA = A.getAAFor<AAReturnedValuesImpl>(*this, *RV);
     if (!RetCSAA) {
@@ -1533,12 +1536,20 @@ struct AAIsDeadFunction : AAIsDead, BooleanState {
     ToBeExploredPaths.insert(&(F.getEntryBlock().front()));
     AssumedLiveBlocks.insert(&(F.getEntryBlock()));
     for (size_t i = 0; i < ToBeExploredPaths.size(); ++i)
-      explorePath(A, ToBeExploredPaths[i]);
+      if (const Instruction *NextNoReturnI =
+              findNextNoReturn(A, ToBeExploredPaths[i]))
+        NoReturnCalls.insert(NextNoReturnI);
   }
 
-  /// Explores new instructions starting from \p I. If instruction is dead, stop
-  /// and return true if it discovered a new instruction.
-  bool explorePath(Attributor &A, Instruction *I);
+  /// Find the next assumed noreturn instruction in the block of \p I starting
+  /// from, thus including, \p I.
+  ///
+  /// The caller is responsible to monitor the ToBeExploredPaths set as new
+  /// instructions discovered in other basic block will be placed in there.
+  ///
+  /// \returns The next assumed noreturn instructions in the block of \p I
+  ///          starting from, thus including, \p I.
+  const Instruction *findNextNoReturn(Attributor &A, const Instruction *I);
 
   const std::string getAsStr() const override {
     return "LiveBBs(" + std::to_string(AssumedLiveBlocks.size()) + "/" +
@@ -1552,18 +1563,31 @@ struct AAIsDeadFunction : AAIsDead, BooleanState {
 
     ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
 
-    for (Instruction *I : NoReturnCalls) {
+    for (const Instruction *NRC : NoReturnCalls) {
+      Instruction *I = const_cast<Instruction *>(NRC);
       BasicBlock *BB = I->getParent();
+      Instruction *SplitPos = I->getNextNode();
 
-      /// Invoke is replaced with a call and unreachable is placed after it.
       if (auto *II = dyn_cast<InvokeInst>(I)) {
-        changeToCall(II);
-        changeToUnreachable(BB->getTerminator(), /* UseLLVMTrap */ false);
-        LLVM_DEBUG(dbgs() << "[AAIsDead] Replaced invoke with call inst\n");
-        continue;
+        /// Invoke is replaced with a call and unreachable is placed after it if
+        /// the callee is nounwind and noreturn. Otherwise, we keep the invoke
+        /// and only place an unreachable in the normal successor.
+        if (Function *Callee = II->getCalledFunction()) {
+          auto *AANoUnw = A.getAAFor<AANoUnwind>(*this, *Callee);
+          if (Callee->hasFnAttribute(Attribute::NoUnwind) ||
+              (AANoUnw && AANoUnw->isAssumedNoUnwind())) {
+            LLVM_DEBUG(dbgs() << "[AAIsDead] Replace invoke with call inst\n");
+            changeToCall(II);
+            changeToUnreachable(BB->getTerminator(), /* UseLLVMTrap */ false);
+            continue;
+          }
+        }
+
+        BB = II->getNormalDest();
+        SplitPos = &BB->front();
       }
 
-      SplitBlock(BB, I->getNextNode());
+      SplitBlock(BB, SplitPos);
       changeToUnreachable(BB->getTerminator(), /* UseLLVMTrap */ false);
       HasChanged = ChangeStatus::CHANGED;
     }
@@ -1575,7 +1599,7 @@ struct AAIsDeadFunction : AAIsDead, BooleanState {
   ChangeStatus updateImpl(Attributor &A) override;
 
   /// See AAIsDead::isAssumedDead(BasicBlock *).
-  bool isAssumedDead(BasicBlock *BB) const override {
+  bool isAssumedDead(const BasicBlock *BB) const override {
     assert(BB->getParent() == &getAnchorScope() &&
            "BB must be in the same anchor scope function.");
 
@@ -1585,12 +1609,12 @@ struct AAIsDeadFunction : AAIsDead, BooleanState {
   }
 
   /// See AAIsDead::isKnownDead(BasicBlock *).
-  bool isKnownDead(BasicBlock *BB) const override {
+  bool isKnownDead(const BasicBlock *BB) const override {
     return getKnown() && isAssumedDead(BB);
   }
 
   /// See AAIsDead::isAssumed(Instruction *I).
-  bool isAssumedDead(Instruction *I) const override {
+  bool isAssumedDead(const Instruction *I) const override {
     assert(I->getParent()->getParent() == &getAnchorScope() &&
            "Instruction must be in the same anchor scope function.");
 
@@ -1603,33 +1627,29 @@ struct AAIsDeadFunction : AAIsDead, BooleanState {
       return true;
 
     // If it is not after a noreturn call, than it is live.
-    if (!isAfterNoReturn(I))
-      return false;
-
-    // Definitely dead.
-    return true;
+    return isAfterNoReturn(I);
   }
 
   /// See AAIsDead::isKnownDead(Instruction *I).
-  bool isKnownDead(Instruction *I) const override {
+  bool isKnownDead(const Instruction *I) const override {
     return getKnown() && isAssumedDead(I);
   }
 
   /// Check if instruction is after noreturn call, in other words, assumed dead.
-  bool isAfterNoReturn(Instruction *I) const;
+  bool isAfterNoReturn(const Instruction *I) const;
 
   /// Collection of to be explored paths.
-  SmallSetVector<Instruction *, 8> ToBeExploredPaths;
+  SmallSetVector<const Instruction *, 8> ToBeExploredPaths;
 
   /// Collection of all assumed live BasicBlocks.
-  DenseSet<BasicBlock *> AssumedLiveBlocks;
+  DenseSet<const BasicBlock *> AssumedLiveBlocks;
 
   /// Collection of calls with noreturn attribute, assumed or knwon.
-  SmallSetVector<Instruction *, 4> NoReturnCalls;
+  SmallSetVector<const Instruction *, 4> NoReturnCalls;
 };
 
-bool AAIsDeadFunction::isAfterNoReturn(Instruction *I) const {
-  Instruction *PrevI = I->getPrevNode();
+bool AAIsDeadFunction::isAfterNoReturn(const Instruction *I) const {
+  const Instruction *PrevI = I->getPrevNode();
   while (PrevI) {
     if (NoReturnCalls.count(PrevI))
       return true;
@@ -1638,75 +1658,77 @@ bool AAIsDeadFunction::isAfterNoReturn(Instruction *I) const {
   return false;
 }
 
-bool AAIsDeadFunction::explorePath(Attributor &A, Instruction *I) {
-  BasicBlock *BB = I->getParent();
+const Instruction *AAIsDeadFunction::findNextNoReturn(Attributor &A,
+                                                      const Instruction *I) {
+  const BasicBlock *BB = I->getParent();
+
+  // TODO: We should have a function that determines if an "edge" is dead.
+  //       Edges could be from an instruction to the next or from a terminator
+  //       to the successor. For now, we need to special case the unwind block
+  //       of InvokeInst below.
 
   while (I) {
     ImmutableCallSite ICS(I);
 
     if (ICS) {
-      auto *NoReturnAA = A.getAAFor<AANoReturn>(*this, *I);
-
-      if (NoReturnAA && NoReturnAA->isAssumedNoReturn()) {
-        if (!NoReturnCalls.insert(I))
-          // If I is already in the NoReturnCalls set, then it stayed noreturn
-          // and we didn't discover any new instructions.
-          return false;
-
-        // Discovered new noreturn call, return true to indicate that I is not
-        // noreturn anymore and should be deleted from NoReturnCalls.
-        return true;
+      // Regarless of the no-return property of an invoke instruction we only
+      // learn that the regular successor is not reachable through this
+      // instruction but the unwind block might still be.
+      if (auto *Invoke = dyn_cast<InvokeInst>(I)) {
+        // Use nounwind to justify the unwind block is dead as well.
+        auto *AANoUnw = A.getAAFor<AANoUnwind>(*this, *Invoke);
+        if (!AANoUnw || !AANoUnw->isAssumedNoUnwind()) {
+          AssumedLiveBlocks.insert(Invoke->getUnwindDest());
+          ToBeExploredPaths.insert(&Invoke->getUnwindDest()->front());
+        }
       }
 
-      if (ICS.hasFnAttr(Attribute::NoReturn)) {
-        if (!NoReturnCalls.insert(I))
-          return false;
-
-        return true;
-      }
+      auto *NoReturnAA = A.getAAFor<AANoReturn>(*this, *I);
+      if (ICS.hasFnAttr(Attribute::NoReturn) ||
+          (NoReturnAA && NoReturnAA->isAssumedNoReturn()))
+        return I;
     }
 
     I = I->getNextNode();
   }
 
   // get new paths (reachable blocks).
-  for (BasicBlock *SuccBB : successors(BB)) {
-    Instruction *Inst = &(SuccBB->front());
+  for (const BasicBlock *SuccBB : successors(BB)) {
     AssumedLiveBlocks.insert(SuccBB);
-    ToBeExploredPaths.insert(Inst);
+    ToBeExploredPaths.insert(&SuccBB->front());
   }
 
-  return true;
+  // No noreturn instruction found.
+  return nullptr;
 }
 
 ChangeStatus AAIsDeadFunction::updateImpl(Attributor &A) {
   // Temporary collection to iterate over existing noreturn instructions. This
   // will alow easier modification of NoReturnCalls collection
-  SmallVector<Instruction *, 8> NoReturnChanged;
+  SmallVector<const Instruction *, 8> NoReturnChanged;
   ChangeStatus Status = ChangeStatus::UNCHANGED;
 
-  for (Instruction *I : NoReturnCalls)
+  for (const Instruction *I : NoReturnCalls)
     NoReturnChanged.push_back(I);
 
-  for (Instruction *I : NoReturnChanged) {
+  for (const Instruction *I : NoReturnChanged) {
     size_t Size = ToBeExploredPaths.size();
 
-    // Still noreturn.
-    if (!explorePath(A, I))
-      continue;
-
-    NoReturnCalls.remove(I);
-
-    // At least one new path.
-    Status = ChangeStatus::CHANGED;
-
-    // No new paths.
-    if (Size == ToBeExploredPaths.size())
-      continue;
+    const Instruction *NextNoReturnI = findNextNoReturn(A, I);
+    if (NextNoReturnI != I) {
+      Status = ChangeStatus::CHANGED;
+      NoReturnCalls.remove(I);
+      if (NextNoReturnI)
+        NoReturnCalls.insert(NextNoReturnI);
+    }
 
-    // explore new paths.
-    while (Size != ToBeExploredPaths.size())
-      explorePath(A, ToBeExploredPaths[Size++]);
+    // Explore new paths.
+    while (Size != ToBeExploredPaths.size()) {
+      Status = ChangeStatus::CHANGED;
+      if (const Instruction *NextNoReturnI =
+              findNextNoReturn(A, ToBeExploredPaths[Size++]))
+        NoReturnCalls.insert(NextNoReturnI);
+    }
   }
 
   LLVM_DEBUG(