[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

3179

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

3348

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 if the edge from \p From BB to \p To BB is assumed dead.
+  /// This is specifically useful in AAReachability.
+  virtual bool isEdgeDead(const BasicBlock *From, const BasicBlock *To) const {
+    return false;
+  }
+
   /// 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;
 
+  bool isEdgeDead(const BasicBlock *From, const BasicBlock *To) const override {
+    return !AssumedLiveEdges.count(std::make_pair(From, To));
+  }
+
   /// 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 edges
+  DenseSet<std::pair<const BasicBlock *, const BasicBlock *>> AssumedLiveEdges;
+
   /// Collection of all assumed live BasicBlocks.
   DenseSet<const BasicBlock *> AssumedLiveBlocks;
 };
 
 static bool
 identifyAliveSuccessors(Attributor &A, const CallBase &CB,
                         AbstractAttribute &AA,
                         SmallVectorImpl<const Instruction *> &AliveSuccessors) {
while (!Worklist.empty()) {
     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 {
+        // record the assumed live edge
+        AssumedLiveEdges.insert(
+            std::make_pair(I->getParent(), AliveSuccessor->getParent()));
         if (assumeLive(A, *AliveSuccessor->getParent()))
           Worklist.push_back(AliveSuccessor);
       }
     }
   }
 
   ToBeExploredFrom = std::move(NewToBeExploredFrom);