[Attributor] Provide an edge-based interface in AAIsDead

[okura]

Summary

This patch produces an edge-based interface in AAIsDead.
By this, we can query a set of basic blocks that are directly reachable from a given basic block.
This is specifically useful for implementation of AAReachability.

[jdoerfert]

llvm/include/llvm/Transforms/IPO/Attributor.h

2654

This will create (and below copy) these vectors all the time, that should be avoided. Can we instead make this a wrapper around succ_begin, e.g., you ask for a successor iterator only for live edges, and you can use it with succ_end. Basically, return an iterator that skips dead edges.

[okura]

• change interface to return an iterator corresponds to the next live successor
• make cache for edges in AAIsDeadFunction

[okura]

• check state validity in getNextAssumedLiveSuccessor

[jdoerfert]

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

3183

Given that a block has usually 1 or 2 successors we should consider a different data structure here. I would initially think a Set<std::pair<BasicBlock*,BasicBlock*>> is adequate, so we collect edges that are assumed live in there.

[okura]

• change whole interface

[jdoerfert]

LGTM, one nit

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

3386

Move this into the else case below. Only if I is a terminator we need to do this and then we know it is an edge. This actually doesn't catch self loops, so move the insert into the else and remove the condition.

[okura]

• fix nit

llvm/include/llvm/Transforms/IPO/Attributor.h

public:
   /// Create an abstract attribute view for the position \p IRP.
   static AAIsDead &createForPosition(const IRPosition &IRP, Attributor &A);
 
   /// Determine if \p F might catch asynchronous exceptions.
   static bool mayCatchAsynchronousExceptions(const Function &F) {
     return F.hasPersonalityFn() && !canSimplifyInvokeNoUnwind(&F);
   }
 
+  /// Return assumed live successors of a given basic block.
+  /// This is specifically useful in AAReachability.
+  virtual SetVector<const BasicBlock *>
+  getAssumedLiveOutEdges(Attributor &A, const BasicBlock &BB) const {
+    return SetVector<const BasicBlock *>(succ_begin(&BB), succ_end(&BB));
+  }
+
   /// See AbstractAttribute::getName()
   const std::string getName() const override { return "AAIsDead"; }
 
   /// See AbstractAttribute::getIdAddr()
   const char *getIdAddr() const override { return &ID; }
 
   /// This function should return true if the type of the \p AA is AAIsDead
   static bool classof(const AbstractAttribute *AA) {

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

for (BasicBlock &BB : F)
       }
 
     return HasChanged;
   }
 
   /// See AbstractAttribute::updateImpl(...).
   ChangeStatus updateImpl(Attributor &A) override;
 
+  /// Returns assumed live successors of a given basic block.
+  /// Note: Even if there is a branch instruction that jumps to the entry node
+  ///       of the given basic block, the returned set will not contain the
+  ///       given basic block.
+  SetVector<const BasicBlock *>
+  getAssumedLiveOutEdges(Attributor &A, const BasicBlock &BB) const override;
+
   /// See AbstractAttribute::trackStatistics()
   void trackStatistics() const override {}
 
   /// Returns true if the function is assumed dead.
   bool isAssumedDead() const override { return false; }
 
   /// See AAIsDead::isKnownDead().
   bool isKnownDead() const override { return false; }
struct AAIsDeadFunction : public AAIsDead {
   /// did assume they do not transfer control to (one of their) successors.
   SmallSetVector<const Instruction *, 8> ToBeExploredFrom;
 
   /// Collection of instructions that are known to not transfer control.
   SmallSetVector<const Instruction *, 8> KnownDeadEnds;
 
   /// Collection of all assumed live BasicBlocks.
   DenseSet<const BasicBlock *> AssumedLiveBlocks;
 };
 
 static bool
 identifyAliveSuccessors(Attributor &A, const CallBase &CB,
-                        AbstractAttribute &AA,
+                        const AbstractAttribute &AA,
                         SmallVectorImpl<const Instruction *> &AliveSuccessors) {
   const IRPosition &IPos = IRPosition::callsite_function(CB);
 
   const auto &NoReturnAA = A.getAndUpdateAAFor<AANoReturn>(
       AA, IPos, /* TrackDependence */ true, DepClassTy::OPTIONAL);
   if (NoReturnAA.isAssumedNoReturn())
     return !NoReturnAA.isKnownNoReturn();
   if (CB.isTerminator())
     AliveSuccessors.push_back(&CB.getSuccessor(0)->front());
   else
     AliveSuccessors.push_back(CB.getNextNode());
   return false;
 }
 
 static bool
 identifyAliveSuccessors(Attributor &A, const InvokeInst &II,
-                        AbstractAttribute &AA,
+                        const AbstractAttribute &AA,
                         SmallVectorImpl<const Instruction *> &AliveSuccessors) {
   bool UsedAssumedInformation =
       identifyAliveSuccessors(A, cast<CallBase>(II), AA, AliveSuccessors);
 
   // First, determine if we can change an invoke to a call assuming the
   // callee is nounwind. This is not possible if the personality of the
   // function allows to catch asynchronous exceptions.
   if (AAIsDeadFunction::mayCatchAsynchronousExceptions(*II.getFunction())) {
     AliveSuccessors.push_back(&II.getUnwindDest()->front());
   } else {
     const IRPosition &IPos = IRPosition::callsite_function(II);
     const auto &AANoUnw = A.getAndUpdateAAFor<AANoUnwind>(
         AA, IPos, /* TrackDependence */ true, DepClassTy::OPTIONAL);
     if (AANoUnw.isAssumedNoUnwind()) {
       UsedAssumedInformation |= !AANoUnw.isKnownNoUnwind();
     } else {
       AliveSuccessors.push_back(&II.getUnwindDest()->front());
     }
   }
   return UsedAssumedInformation;
 }
 
 static bool
 identifyAliveSuccessors(Attributor &A, const BranchInst &BI,
-                        AbstractAttribute &AA,
+                        const AbstractAttribute &AA,
                         SmallVectorImpl<const Instruction *> &AliveSuccessors) {
   bool UsedAssumedInformation = false;
   if (BI.getNumSuccessors() == 1) {
     AliveSuccessors.push_back(&BI.getSuccessor(0)->front());
   } else {
     Optional<ConstantInt *> CI = getAssumedConstantInt(
         A, *BI.getCondition(), AA, UsedAssumedInformation);
     if (!CI.hasValue()) {
       // No value yet, assume both edges are dead.
     } else if (CI.getValue()) {
       const BasicBlock *SuccBB =
           BI.getSuccessor(1 - CI.getValue()->getZExtValue());
       AliveSuccessors.push_back(&SuccBB->front());
     } else {
       AliveSuccessors.push_back(&BI.getSuccessor(0)->front());
       AliveSuccessors.push_back(&BI.getSuccessor(1)->front());
       UsedAssumedInformation = false;
     }
   }
   return UsedAssumedInformation;
 }
 
 static bool
 identifyAliveSuccessors(Attributor &A, const SwitchInst &SI,
-                        AbstractAttribute &AA,
+                        const AbstractAttribute &AA,
                         SmallVectorImpl<const Instruction *> &AliveSuccessors) {
   bool UsedAssumedInformation = false;
   Optional<ConstantInt *> CI =
       getAssumedConstantInt(A, *SI.getCondition(), AA, UsedAssumedInformation);
   if (!CI.hasValue()) {
     // No value yet, assume all edges are dead.
   } else if (CI.getValue()) {
     for (auto &CaseIt : SI.cases()) {
       if (CaseIt.getCaseValue() == CI.getValue()) {
         AliveSuccessors.push_back(&CaseIt.getCaseSuccessor()->front());
         return UsedAssumedInformation;
       }
     }
     AliveSuccessors.push_back(&SI.getDefaultDest()->front());
     return UsedAssumedInformation;
   } else {
     for (const BasicBlock *SuccBB : successors(SI.getParent()))
       AliveSuccessors.push_back(&SuccBB->front());
   }
   return UsedAssumedInformation;
 }
 
+static bool identifyAliveSuccessorsForInstruction(
+    Attributor &A, const Instruction &I, const AbstractAttribute &AA,
+    SmallVectorImpl<const Instruction *> &AliveSuccessors) {
+  bool UsedAssumedInformation = false;
+  switch (I.getOpcode()) {
+  // TODO: look for (assumed) UB to backwards propagate "deadness".
+  default:
+    if (I.isTerminator()) {
+      for (const BasicBlock *SuccBB : successors(I.getParent()))
+        AliveSuccessors.push_back(&SuccBB->front());
+    } else {
+      AliveSuccessors.push_back(I.getNextNode());
+    }
+    break;
+  case Instruction::Call:
+    UsedAssumedInformation =
+        identifyAliveSuccessors(A, cast<CallInst>(I), AA, AliveSuccessors);
+    break;
+  case Instruction::Invoke:
+    UsedAssumedInformation =
+        identifyAliveSuccessors(A, cast<InvokeInst>(I), AA, AliveSuccessors);
+    break;
+  case Instruction::Br:
+    UsedAssumedInformation =
+        identifyAliveSuccessors(A, cast<BranchInst>(I), AA, AliveSuccessors);
+    break;
+  case Instruction::Switch:
+    UsedAssumedInformation =
+        identifyAliveSuccessors(A, cast<SwitchInst>(I), AA, AliveSuccessors);
+    break;
+  }
+  return UsedAssumedInformation;
+}
+
+SetVector<const BasicBlock *>
+AAIsDeadFunction::getAssumedLiveOutEdges(Attributor &A,
+                                         const BasicBlock &BB) const {
+  // Explore from the entry node of the given basic block.
+  SmallVector<const Instruction *, 8> Worklist;
+  DenseSet<const Instruction *> ExploredInsts;
+  SmallVector<const Instruction *, 8> AliveSuccessors;
+  SetVector<const BasicBlock *> AliveSuccessorBBs;
+
+  Worklist.push_back(&BB.front());
+  ExploredInsts.insert(&BB.front());
+
+  while (!Worklist.empty()) {
+    const Instruction *I = Worklist.pop_back_val();
+    AliveSuccessors.clear();
+    identifyAliveSuccessorsForInstruction(A, *I, *this, AliveSuccessors);
+
+    for (const Instruction *AliveSuccessor : AliveSuccessors) {
+      const BasicBlock *AliveSuccessorBB = AliveSuccessor->getParent();
+      // If a successor is outside of the basic block, add it to the assumed
+      // live successor set. Otherwise, If we have not explore the successor
+      // yet, add it to the worklist.
+      if (AliveSuccessorBB != &BB)
+        AliveSuccessorBBs.insert(AliveSuccessorBB);
+      else if (!ExploredInsts.count(AliveSuccessor)) {
+        Worklist.push_back(AliveSuccessor);
+        ExploredInsts.insert(AliveSuccessor);
+      }
+    }
+  }
+
+  return AliveSuccessorBBs;
+}
+
 ChangeStatus AAIsDeadFunction::updateImpl(Attributor &A) {
   ChangeStatus Change = ChangeStatus::UNCHANGED;
 
   LLVM_DEBUG(dbgs() << "[AAIsDead] Live [" << AssumedLiveBlocks.size() << "/"
                     << getAnchorScope()->size() << "] BBs and "
                     << ToBeExploredFrom.size() << " exploration points and "
                     << KnownDeadEnds.size() << " known dead ends\n");
 
   // Copy and clear the list of instructions we need to explore from. It is
   // refilled with instructions the next update has to look at.
   SmallVector<const Instruction *, 8> Worklist(ToBeExploredFrom.begin(),
                                                ToBeExploredFrom.end());
   decltype(ToBeExploredFrom) NewToBeExploredFrom;
 
   SmallVector<const Instruction *, 8> AliveSuccessors;
   while (!Worklist.empty()) {
     const Instruction *I = Worklist.pop_back_val();
     LLVM_DEBUG(dbgs() << "[AAIsDead] Exploration inst: " << *I << "\n");
 
     AliveSuccessors.clear();
 
-    bool UsedAssumedInformation = false;
-    switch (I->getOpcode()) {
-    // TODO: look for (assumed) UB to backwards propagate "deadness".
-    default:
-      if (I->isTerminator()) {
-        for (const BasicBlock *SuccBB : successors(I->getParent()))
-          AliveSuccessors.push_back(&SuccBB->front());
-      } else {
-        AliveSuccessors.push_back(I->getNextNode());
-      }
-      break;
-    case Instruction::Call:
-      UsedAssumedInformation = identifyAliveSuccessors(A, cast<CallInst>(*I),
-                                                       *this, AliveSuccessors);
-      break;
-    case Instruction::Invoke:
-      UsedAssumedInformation = identifyAliveSuccessors(A, cast<InvokeInst>(*I),
-                                                       *this, AliveSuccessors);
-      break;
-    case Instruction::Br:
-      UsedAssumedInformation = identifyAliveSuccessors(A, cast<BranchInst>(*I),
-                                                       *this, AliveSuccessors);
-      break;
-    case Instruction::Switch:
-      UsedAssumedInformation = identifyAliveSuccessors(A, cast<SwitchInst>(*I),
-                                                       *this, AliveSuccessors);
-      break;
-    }
+    bool UsedAssumedInformation =
+        identifyAliveSuccessorsForInstruction(A, *I, *this, AliveSuccessors);
 
     if (UsedAssumedInformation) {
       NewToBeExploredFrom.insert(I);
     } else {
       Change = ChangeStatus::CHANGED;
       if (AliveSuccessors.empty() ||
           (I->isTerminator() && AliveSuccessors.size() < I->getNumSuccessors()))
         KnownDeadEnds.insert(I);
     }
 
     LLVM_DEBUG(dbgs() << "[AAIsDead] #AliveSuccessors: "
                       << AliveSuccessors.size() << " UsedAssumedInformation: "
                       << UsedAssumedInformation << "\n");
 
     for (const Instruction *AliveSuccessor : AliveSuccessors) {
       if (!I->isTerminator()) {
         assert(AliveSuccessors.size() == 1 &&
                "Non-terminator expected to have a single successor!");
         Worklist.push_back(AliveSuccessor);
       } else {
         if (assumeLive(A, *AliveSuccessor->getParent()))
           Worklist.push_back(AliveSuccessor);
       }
     }
   }
 
   ToBeExploredFrom = std::move(NewToBeExploredFrom);