[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

3192

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

3365

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 const_succ_iterator
+  getNextAssumedLiveSuccessor(const_succ_iterator It) const {
+    return It++;
+  }
+
   /// 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 an iterator that corresponds to the next assumed live successor.
+  /// If there is not such successor, returns the end iterator.
+  const_succ_iterator
+  getNextAssumedLiveSuccessor(const_succ_iterator It) const override {
+    const BasicBlock *BB = It.getSource();
+    while (++It != succ_end(BB)) {
+      auto AssumedLiveSuccessorsIt = AssumedLiveSuccessors.find(BB);
+      if (AssumedLiveSuccessorsIt == AssumedLiveSuccessors.end())
+        return It;
+      const BasicBlock *SuccBB = *It;
+      if (AssumedLiveSuccessorsIt->second->count(SuccBB))
+        return It;
+    }
+    return It;
+  }
+
   /// 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 {
 
   /// Collection of instructions that need to be explored again, e.g., we
   /// 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 successors for each block
+  DenseMap<const BasicBlock *, std::unique_ptr<DenseSet<const BasicBlock *>>>
+      AssumedLiveSuccessors;
+
   /// Collection of all assumed live BasicBlocks.
   DenseSet<const BasicBlock *> AssumedLiveBlocks;
 };
 
 static bool
 identifyAliveSuccessors(Attributor &A, const CallBase &CB,
                         AbstractAttribute &AA,
                         SmallVectorImpl<const Instruction *> &AliveSuccessors) {
if (!CI.hasValue()) {
     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, 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;
+}
+
 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 successor belongs to different basic block, record the edge.
+      if (I->getParent() != AliveSuccessor->getParent()) {
+        auto AssumedLiveSuccessorsInsertResult =
+            AssumedLiveSuccessors.insert(std::make_pair(
+                I->getParent(), std::unique_ptr<DenseSet<const BasicBlock *>>(
+                                    new DenseSet<const BasicBlock *>())));
+        AssumedLiveSuccessorsInsertResult.first->second->insert(
+            AliveSuccessor->getParent());
+      }
+
       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);
       }