[MustExecute] Add backward exploration for must-be-executed-context

Summary:
As mentioned in D71974, it is useful for must-be-executed-context to explore CFG backwardly.
This patch is ported from parts of D64975. We use a dominator tree to find the previous context if
a dominator tree is available.

Reviewers: jdoerfert, hfinkel, baziotis, sstefan1

Reviewed By: jdoerfert

Subscribers: hiraditya, llvm-commits

Tags: #llvm

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

Author: uenoku

llvm/include/llvm/Analysis/MustExecute.h

@@ -178,6 +178,12 @@ class ICFLoopSafetyInfo: public LoopSafetyInfo {
 
 struct MustBeExecutedContextExplorer;
 
+/// Enum that allows us to spell out the direction.
+enum class ExplorationDirection {
+  BACKWARD = 0,
+  FORWARD = 1,
+};
+
 /// Must be executed iterators visit stretches of instructions that are
 /// guaranteed to be executed together, potentially with other instruction
 /// executed in-between.
@@ -282,16 +288,18 @@ struct MustBeExecutedIterator {
 
   MustBeExecutedIterator(const MustBeExecutedIterator &Other)
       : Visited(Other.Visited), Explorer(Other.Explorer),
-        CurInst(Other.CurInst) {}
+        CurInst(Other.CurInst), Head(Other.Head), Tail(Other.Tail) {}
 
   MustBeExecutedIterator(MustBeExecutedIterator &&Other)
       : Visited(std::move(Other.Visited)), Explorer(Other.Explorer),
-        CurInst(Other.CurInst) {}
+        CurInst(Other.CurInst), Head(Other.Head), Tail(Other.Tail) {}
 
   MustBeExecutedIterator &operator=(MustBeExecutedIterator &&Other) {
     if (this != &Other) {
       std::swap(Visited, Other.Visited);
       std::swap(CurInst, Other.CurInst);
+      std::swap(Head, Other.Head);
+      std::swap(Tail, Other.Tail);
     }
     return *this;
   }
@@ -315,7 +323,7 @@ struct MustBeExecutedIterator {
   /// Equality and inequality operators. Note that we ignore the history here.
   ///{
   bool operator==(const MustBeExecutedIterator &Other) const {
-    return CurInst == Other.CurInst;
+    return CurInst == Other.CurInst && Head == Other.Head && Tail == Other.Tail;
   }
 
   bool operator!=(const MustBeExecutedIterator &Other) const {
@@ -328,17 +336,24 @@ struct MustBeExecutedIterator {
   const Instruction *getCurrentInst() const { return CurInst; }
 
   /// Return true if \p I was encountered by this iterator already.
-  bool count(const Instruction *I) const { return Visited.count(I); }
+  bool count(const Instruction *I) const {
+    return Visited.count({I, ExplorationDirection::FORWARD}) ||
+           Visited.count({I, ExplorationDirection::BACKWARD});
+  }
 
 private:
-  using VisitedSetTy = DenseSet<const Instruction *>;
+  using VisitedSetTy =
+      DenseSet<PointerIntPair<const Instruction *, 1, ExplorationDirection>>;
 
   /// Private constructors.
   MustBeExecutedIterator(ExplorerTy &Explorer, const Instruction *I);
 
   /// Reset the iterator to its initial state pointing at \p I.
   void reset(const Instruction *I);
 
+  /// Reset the iterator to point at \p I, keep cached state.
+  void resetInstruction(const Instruction *I);
+
   /// Try to advance one of the underlying positions (Head or Tail).
   ///
   /// \return The next instruction in the must be executed context, or nullptr
@@ -357,6 +372,11 @@ struct MustBeExecutedIterator {
   /// initially the program point itself.
   const Instruction *CurInst;
 
+  /// Two positions that mark the program points where this iterator will look
+  /// for the next instruction. Note that the current instruction is either the
+  /// one pointed to by Head, Tail, or both.
+  const Instruction *Head, *Tail;
+
   friend struct MustBeExecutedContextExplorer;
 };
 
@@ -379,14 +399,24 @@ struct MustBeExecutedContextExplorer {
   /// \param ExploreInterBlock    Flag to indicate if instructions in blocks
   ///                             other than the parent of PP should be
   ///                             explored.
+  /// \param ExploreCFGForward    Flag to indicate if instructions located after
+  ///                             PP in the CFG, e.g., post-dominating PP,
+  ///                             should be explored.
+  /// \param ExploreCFGBackward   Flag to indicate if instructions located
+  ///                             before PP in the CFG, e.g., dominating PP,
+  ///                             should be explored.
   MustBeExecutedContextExplorer(
-      bool ExploreInterBlock,
+      bool ExploreInterBlock, bool ExploreCFGForward, bool ExploreCFGBackward,
       GetterTy<const LoopInfo> LIGetter =
           [](const Function &) { return nullptr; },
+      GetterTy<const DominatorTree> DTGetter =
+          [](const Function &) { return nullptr; },
       GetterTy<const PostDominatorTree> PDTGetter =
           [](const Function &) { return nullptr; })
-      : ExploreInterBlock(ExploreInterBlock), LIGetter(LIGetter),
-        PDTGetter(PDTGetter), EndIterator(*this, nullptr) {}
+      : ExploreInterBlock(ExploreInterBlock),
+        ExploreCFGForward(ExploreCFGForward),
+        ExploreCFGBackward(ExploreCFGBackward), LIGetter(LIGetter),
+        DTGetter(DTGetter), PDTGetter(PDTGetter), EndIterator(*this, nullptr) {}
 
   /// Clean up the dynamically allocated iterators.
   ~MustBeExecutedContextExplorer() {
@@ -464,21 +494,36 @@ struct MustBeExecutedContextExplorer {
   const Instruction *
   getMustBeExecutedNextInstruction(MustBeExecutedIterator &It,
                                    const Instruction *PP);
+  /// Return the previous instr. that is guaranteed to be executed before \p PP.
+  ///
+  /// \param It              The iterator that is used to traverse the must be
+  ///                        executed context.
+  /// \param PP              The program point for which the previous instr.
+  ///                        that is guaranteed to execute is determined.
+  const Instruction *
+  getMustBeExecutedPrevInstruction(MustBeExecutedIterator &It,
+                                   const Instruction *PP);
 
   /// Find the next join point from \p InitBB in forward direction.
   const BasicBlock *findForwardJoinPoint(const BasicBlock *InitBB);
 
+  /// Find the next join point from \p InitBB in backward direction.
+  const BasicBlock *findBackwardJoinPoint(const BasicBlock *InitBB);
+
   /// Parameter that limit the performed exploration. See the constructor for
   /// their meaning.
   ///{
   const bool ExploreInterBlock;
+  const bool ExploreCFGForward;
+  const bool ExploreCFGBackward;
   ///}
 
 private:
   /// Getters for common CFG analyses: LoopInfo, DominatorTree, and
   /// PostDominatorTree.
   ///{
   GetterTy<const LoopInfo> LIGetter;
+  GetterTy<const DominatorTree> DTGetter;
   GetterTy<const PostDominatorTree> PDTGetter;
   ///}
 

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

@@ -569,8 +569,10 @@ struct AnalysisGetter {
 struct InformationCache {
   InformationCache(const Module &M, AnalysisGetter &AG,
                    SetVector<Function *> *CGSCC)
-      : DL(M.getDataLayout()), Explorer(/* ExploreInterBlock */ true), AG(AG),
-        CGSCC(CGSCC) {}
+      : DL(M.getDataLayout()),
+        Explorer(/* ExploreInterBlock */ true, /* ExploreCFGForward */ true,
+                 /* ExploreCFGBackward */ true),
+        AG(AG), CGSCC(CGSCC) {}
 
   /// A map type from opcodes to instructions with this opcode.
   using OpcodeInstMapTy = DenseMap<unsigned, SmallVector<Instruction *, 32>>;

llvm/lib/Analysis/MustExecute.cpp

@@ -368,12 +368,21 @@ bool MustBeExecutedContextPrinter::runOnModule(Module &M) {
     LIs.push_back(LI);
     return LI;
   };
+  GetterTy<DominatorTree> DTGetter = [&](const Function &F) {
+    DominatorTree *DT = new DominatorTree(const_cast<Function &>(F));
+    DTs.push_back(DT);
+    return DT;
+  };
   GetterTy<PostDominatorTree> PDTGetter = [&](const Function &F) {
     PostDominatorTree *PDT = new PostDominatorTree(const_cast<Function &>(F));
     PDTs.push_back(PDT);
     return PDT;
   };
-  MustBeExecutedContextExplorer Explorer(true, LIGetter, PDTGetter);
+  MustBeExecutedContextExplorer Explorer(
+      /* ExploreInterBlock */ true,
+      /* ExploreCFGForward */ true,
+      /* ExploreCFGBackward */ true, LIGetter, DTGetter, PDTGetter);
+
   for (Function &F : M) {
     for (Instruction &I : instructions(F)) {
       dbgs() << "-- Explore context of: " << I << "\n";
@@ -632,6 +641,72 @@ MustBeExecutedContextExplorer::findForwardJoinPoint(const BasicBlock *InitBB) {
   LLVM_DEBUG(dbgs() << "\tJoin block: " << JoinBB->getName() << "\n");
   return JoinBB;
 }
+const BasicBlock *
+MustBeExecutedContextExplorer::findBackwardJoinPoint(const BasicBlock *InitBB) {
+  const LoopInfo *LI = LIGetter(*InitBB->getParent());
+  const DominatorTree *DT = DTGetter(*InitBB->getParent());
+  LLVM_DEBUG(dbgs() << "\tFind backward join point for " << InitBB->getName()
+                    << (LI ? " [LI]" : "") << (DT ? " [DT]" : ""));
+
+  // Try to determine a join block through the help of the dominance tree. If no
+  // tree was provided, we perform simple pattern matching for one block
+  // conditionals only.
+  if (DT)
+    if (const auto *InitNode = DT->getNode(InitBB))
+      if (const auto *IDomNode = InitNode->getIDom())
+        return IDomNode->getBlock();
+
+  const Loop *L = LI ? LI->getLoopFor(InitBB) : nullptr;
+  const BasicBlock *HeaderBB = L ? L->getHeader() : nullptr;
+
+  // Determine the predecessor blocks but ignore backedges.
+  SmallVector<const BasicBlock *, 8> Worklist;
+  for (const BasicBlock *PredBB : predecessors(InitBB)) {
+    bool IsBackedge =
+        (PredBB == InitBB) || (HeaderBB == InitBB && L->contains(PredBB));
+    // Loop backedges are ignored in backwards propagation: control has to come
+    // from somewhere.
+    if (!IsBackedge)
+      Worklist.push_back(PredBB);
+  }
+
+  // If there are no other predecessor blocks, there is no join point.
+  if (Worklist.empty())
+    return nullptr;
+
+  // If there is one predecessor block, it is the join point.
+  if (Worklist.size() == 1)
+    return Worklist[0];
+
+  const BasicBlock *JoinBB = nullptr;
+  if (Worklist.size() == 2) {
+    const BasicBlock *Pred0 = Worklist[0];
+    const BasicBlock *Pred1 = Worklist[1];
+    const BasicBlock *Pred0UniquePred = Pred0->getUniquePredecessor();
+    const BasicBlock *Pred1UniquePred = Pred1->getUniquePredecessor();
+    if (Pred0 == Pred1UniquePred) {
+      // InitBB <-          Pred0 = JoinBB
+      // InitBB <- Pred1 <- Pred0 = JoinBB
+      JoinBB = Pred0;
+    } else if (Pred1 == Pred0UniquePred) {
+      // InitBB <- Pred0 <- Pred1 = JoinBB
+      // InitBB <-          Pred1 = JoinBB
+      JoinBB = Pred1;
+    } else if (Pred0UniquePred == Pred1UniquePred) {
+      // InitBB <- Pred0 <- JoinBB
+      // InitBB <- Pred1 <- JoinBB
+      JoinBB = Pred0UniquePred;
+    }
+  }
+
+  if (!JoinBB && L)
+    JoinBB = L->getHeader();
+
+  // In backwards direction there is no need to show termination of previous
+  // instructions. If they do not terminate, the code afterward is dead, making
+  // any information/transformation correct anyway.
+  return JoinBB;
+}
 
 const Instruction *
 MustBeExecutedContextExplorer::getMustBeExecutedNextInstruction(
@@ -690,23 +765,79 @@ MustBeExecutedContextExplorer::getMustBeExecutedNextInstruction(
   return nullptr;
 }
 
+const Instruction *
+MustBeExecutedContextExplorer::getMustBeExecutedPrevInstruction(
+    MustBeExecutedIterator &It, const Instruction *PP) {
+  if (!PP)
+    return PP;
+
+  bool IsFirst = !(PP->getPrevNode());
+  LLVM_DEBUG(dbgs() << "Find next instruction for " << *PP
+                    << (IsFirst ? " [IsFirst]" : "") << "\n");
+
+  // If we explore only inside a given basic block we stop at the first
+  // instruction.
+  if (!ExploreInterBlock && IsFirst) {
+    LLVM_DEBUG(dbgs() << "\tReached block front in intra-block mode, done\n");
+    return nullptr;
+  }
+
+  // The block and function that contains the current position.
+  const BasicBlock *PPBlock = PP->getParent();
+
+  // If we are inside a block we know what instruction was executed before, the
+  // previous one.
+  if (!IsFirst) {
+    const Instruction *PrevPP = PP->getPrevNode();
+    LLVM_DEBUG(
+        dbgs() << "\tIntermediate instruction, continue with previous\n");
+    // We did not enter a callee so we simply return the previous instruction.
+    return PrevPP;
+  }
+
+  // Finally, we have to handle the case where the program point is the first in
+  // a block but not in the function. We use the findBackwardJoinPoint helper
+  // function with information about the function and helper analyses, if
+  // available.
+  if (const BasicBlock *JoinBB = findBackwardJoinPoint(PPBlock))
+    return &JoinBB->back();
+
+  LLVM_DEBUG(dbgs() << "\tNo join point found\n");
+  return nullptr;
+}
+
 MustBeExecutedIterator::MustBeExecutedIterator(
     MustBeExecutedContextExplorer &Explorer, const Instruction *I)
     : Explorer(Explorer), CurInst(I) {
   reset(I);
 }
 
 void MustBeExecutedIterator::reset(const Instruction *I) {
-  CurInst = I;
   Visited.clear();
-  Visited.insert(I);
+  resetInstruction(I);
+}
+
+void MustBeExecutedIterator::resetInstruction(const Instruction *I) {
+  CurInst = I;
+  Head = Tail = nullptr;
+  Visited.insert({I, ExplorationDirection::FORWARD});
+  Visited.insert({I, ExplorationDirection::BACKWARD});
+  if (Explorer.ExploreCFGForward)
+    Head = I;
+  if (Explorer.ExploreCFGBackward)
+    Tail = I;
 }
 
 const Instruction *MustBeExecutedIterator::advance() {
   assert(CurInst && "Cannot advance an end iterator!");
-  const Instruction *Next =
-      Explorer.getMustBeExecutedNextInstruction(*this, CurInst);
-  if (Next && !Visited.insert(Next).second)
-    Next = nullptr;
-  return Next;
+  Head = Explorer.getMustBeExecutedNextInstruction(*this, Head);
+  if (Head && Visited.insert({Head, ExplorationDirection ::FORWARD}).second)
+    return Head;
+  Head = nullptr;
+
+  Tail = Explorer.getMustBeExecutedPrevInstruction(*this, Tail);
+  if (Tail && Visited.insert({Tail, ExplorationDirection ::BACKWARD}).second)
+    return Tail;
+  Tail = nullptr;
+  return nullptr;
 }