[MustExecute] Forward iterate over conditional branches

[jdoerfert]

Summary

If a conditional branch is encountered we can try to find a join block
where the execution is known to continue. This means finding a suitable
block, e.g., the immediate post dominator of the conditional branch, and
proofing control will always reach that block.

This patch implements different techniques that work with and without
provided analysis.

[uenoku]

Basically it looks good to me.
I think we need a test for a loop(never-endless + propagation from exit block).
Something like,

;int g(int*);
;void h(int*);
;
;int f(int *u, int n){
;  for(int i = 0;i<n;i++){
;    h(u);
;  }
;  return g(nonnull u);
;}

define i32 @f(i32* %a, i32 %b) local_unnamed_addr {
  %3 = icmp eq i32 %b, 0
  br i1 %3, label %4, label %6

; <label>:4:                                      ; preds = %6, %2
  %5 = tail call i32 @g(i32* nonnull %a)
  ret i32 %5

; <label>:6:                                      ; preds = %2, %6
  %7 = phi i32 [ %8, %6 ], [ 0, %2 ]
  tail call void @h(i32* %a)
  %8 = add nuw i32 %7, 1
  %9 = icmp eq i32 %8, %b
  br i1 %9, label %4, label %6
}

declare void @h(i32*) willreturn nounwind 
declare i32 @g(i32*) willreturn nounwind

[jdoerfert]

I'll add that test (in MustExecute and in the Attributor tests) but it will not be deduced yet. We need better loop support first.

[jdoerfert]

Do I assume correctly that this is accepted with this additional test?

[uenoku]

In D68933#1725277, @jdoerfert wrote:

Do I assume correctly that this is accepted with this additional test?

Yes

llvm/include/llvm/Analysis/MustExecute.h

 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instruction.h"
 
 namespace llvm {
 
 class Instruction;
 class DominatorTree;
+class PostDominatorTree;
 class Loop;
 
 /// Captures loop safety information.
 /// It keep information for loop blocks may throw exception or otherwise
 /// exit abnormaly on any iteration of the loop which might actually execute
 /// at runtime.  The primary way to consume this infromation is via
 /// isGuaranteedToExecute below, but some callers bailout or fallback to
 /// alternate reasoning if a loop contains any implicit control flow.
struct MustBeExecutedContextExplorer {
 
   /// In the description of the parameters we use PP to denote a program point
   /// for which the must be executed context is explored, or put differently,
   /// for which the MustBeExecutedIterator is created.
   ///
   /// \param ExploreInterBlock    Flag to indicate if instructions in blocks
   ///                             other than the parent of PP should be
   ///                             explored.
-  MustBeExecutedContextExplorer(bool ExploreInterBlock)
-      : ExploreInterBlock(ExploreInterBlock), EndIterator(*this, nullptr),
+  MustBeExecutedContextExplorer(bool ExploreInterBlock,
+                                const LoopInfo *LI = nullptr,
+                                const PostDominatorTree *PDT = nullptr)
+      : ExploreInterBlock(ExploreInterBlock), LI(LI), PDT(PDT),
+        EndIterator(*this, nullptr),
         IsGuaranteedToTransferFn([](const Instruction *I) {
           return isGuaranteedToTransferExecutionToSuccessor(I);
         }) {}
 
   /// Clean up the dynamically allocated iterators.
   ~MustBeExecutedContextExplorer() {
     DeleteContainerSeconds(InstructionIteratorMap);
   }
struct MustBeExecutedContextExplorer {
   /// \param It              The iterator that is used to traverse the must be
   ///                        executed context.
   /// \param PP              The program point for which the next instruction
   ///                        that is guaranteed to execute is determined.
   const Instruction *
   getMustBeExecutedNextInstruction(MustBeExecutedIterator &It,
                                    const Instruction *PP);
 
+  /// Find the next join point from \p InitBB in forward direction.
+  const BasicBlock *findForwardJoinPoint(const BasicBlock *InitBB);
+
   /// Parameter that limit the performed exploration. See the constructor for
   /// their meaning.
   ///{
   const bool ExploreInterBlock;
   ///}
 
   /// Replace the function that determines if an instruction "terminates" with
   /// \p IsGuaranteedToTransferFn.
   void setIsGuaranteedToTransferFn(
       IsGuaranteedToTransferFnTy IsGuaranteedToTransferFn) {
     this->IsGuaranteedToTransferFn = IsGuaranteedToTransferFn;
   }
 
 private:
+  /// CFG analyses
+  ///
+  ///{
+  const LoopInfo* LI;
+  const PostDominatorTree* PDT;
+  ///}
+
+  /// Map to cache isGuaranteedToTransferExecutionToSuccessor results.
+  DenseMap<const BasicBlock *, Optional<bool>> BlockTransferMap;
+
+  /// Map to cache containsIrreducibleCFG results.
+  DenseMap<const Function*, Optional<bool>> IrreducibleControlMap;
+
   /// Map from instructions to associated must be executed iterators.
   DenseMap<const Instruction *, MustBeExecutedIterator *>
       InstructionIteratorMap;
 
   /// A unique end iterator.
   MustBeExecutedIterator EndIterator;
 
   /// The callback to determines if a function *will* tranfer execution to its
   /// successor. By default isGuaranteedToTransferExecutionToSuccessor(const
   /// Instruction *) in ValueTracking.h.
   IsGuaranteedToTransferFnTy IsGuaranteedToTransferFn;
 };
 
 } // namespace llvm
 
 #endif

llvm/lib/Analysis/MustExecute.cpp

 //===- MustExecute.cpp - Printer for isGuaranteedToExecute ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/MustExecute.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/Passes.h"
+#include "llvm/Analysis/PostDominators.h"
 #include "llvm/IR/AssemblyAnnotationWriter.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FormattedStream.h"
 #include "llvm/Support/raw_ostream.h"
bool MustExecutePrinter::runOnFunction(Function &F) {
   auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
 
   MustExecuteAnnotatedWriter Writer(F, DT, LI);
   F.print(dbgs(), &Writer);
 
   return false;
 }
 
+/// Return true if \p L might be an endless loop.
+static bool maybeEndlessLoop(const Loop &L) {
+  if (L.getHeader()->getParent()->hasFnAttribute(Attribute::WillReturn))
+    return false;
+  // TODO: Actually try to prove it is not.
+  // TODO: If maybeEndlessLoop is going to be expensive, cache it.
+  return true;
+}
+
+static bool mayContainIrreducibleControl(const Function &F, const LoopInfo *LI) {
+  if (!LI)
+    return false;
+  using RPOTraversal = ReversePostOrderTraversal<const Function *>;
+  RPOTraversal FuncRPOT(&F);
+  return !containsIrreducibleCFG<const BasicBlock *, const RPOTraversal,
+                                 const LoopInfo>(FuncRPOT, *LI);
+}
+
+/// Lookup \p Key in \p Map and return the result, potentially after
+/// initializing the optional through \p Fn(\p args).
+template <typename K, typename V, typename FnTy, typename... ArgsTy>
+static V getOrCreateCachedOptional(K Key, DenseMap<K, Optional<V>> &Map,
+                                   FnTy &&Fn, ArgsTy&&... args) {
+  Optional<V> &OptVal = Map[Key];
+  if (!OptVal.hasValue())
+    OptVal = Fn(std::forward<ArgsTy>(args)...);
+  return OptVal.getValue();
+}
+
+const BasicBlock *
+MustBeExecutedContextExplorer::findForwardJoinPoint(const BasicBlock *InitBB) {
+  LLVM_DEBUG(dbgs() << "\tFind forward join point for " << InitBB->getName()
+                    << (LI ? " [LI]" : "") << (PDT ? " [PDT]" : ""));
+
+  const Function &F = *InitBB->getParent();
+  const Loop *L = LI ? LI->getLoopFor(InitBB) : nullptr;
+  const BasicBlock *HeaderBB = L ? L->getHeader() : InitBB;
+  bool WillReturnAndNoThrow = (F.hasFnAttribute(Attribute::WillReturn) ||
+                               (L && !maybeEndlessLoop(*L))) &&
+                              F.doesNotThrow();
+  LLVM_DEBUG(dbgs() << (L ? " [in loop]" : "")
+                    << (WillReturnAndNoThrow ? " [WillReturn] [NoUnwind]" : "")
+                    << "\n");
+
+  // Determine the adjacent blocks in the given direction but exclude (self)
+  // loops under certain circumstances.
+  SmallVector<const BasicBlock *, 8> Worklist;
+  for (const BasicBlock *SuccBB : successors(InitBB)) {
+    bool IsLatch = SuccBB == HeaderBB;
+    // Loop latches are ignored in forward propagation if the loop cannot be
+    // endless and may not throw: control has to go somewhere.
+    if (!WillReturnAndNoThrow || !IsLatch)
+      Worklist.push_back(SuccBB);
+  }
+  LLVM_DEBUG(dbgs() << "\t\t#Worklist: " << Worklist.size() << "\n");
+
+  // If there are no other adjacent blocks, there is no join point.
+  if (Worklist.empty())
+    return nullptr;
+
+  // If there is one adjacent block, it is the join point.
+  if (Worklist.size() == 1)
+    return Worklist[0];
+
+  // Try to determine a join block through the help of the post-dominance
+  // tree. If no tree was provided, we perform simple pattern matching for one
+  // block conditionals only.
+  const BasicBlock *JoinBB = nullptr;
+  if (PDT)
+    if (const auto *InitNode = PDT->getNode(InitBB))
+      if (const auto *IDomNode = InitNode->getIDom())
+        JoinBB = IDomNode->getBlock();
+
+  if (!JoinBB && Worklist.size() == 2) {
+    const BasicBlock *Succ0 = Worklist[0];
+    const BasicBlock *Succ1 = Worklist[1];
+    const BasicBlock *Succ0UniqueSucc = Succ0->getUniqueSuccessor();
+    const BasicBlock *Succ1UniqueSucc = Succ1->getUniqueSuccessor();
+    if (Succ0 == Succ1UniqueSucc) {
+      // InitBB ->          Succ0 = JoinBB
+      // InitBB -> Succ1 -> Succ0 = JoinBB
+      JoinBB = Succ0;
+    } else if (Succ1 == Succ0UniqueSucc) {
+      // InitBB -> Succ0 -> Succ1 = JoinBB
+      // InitBB ->          Succ1 = JoinBB
+      JoinBB = Succ1;
+    } else if (Succ0UniqueSucc == Succ1UniqueSucc) {
+      // InitBB -> Succ0 -> JoinBB
+      // InitBB -> Succ1 -> JoinBB
+      JoinBB = Succ0UniqueSucc;
+    }
+  }
+
+  if (!JoinBB && L)
+    JoinBB = L->getUniqueExitBlock();
+
+  if (!JoinBB)
+    return nullptr;
+
+  LLVM_DEBUG(dbgs() << "\t\tJoin block candidate: " << JoinBB->getName() << "\n");
+
+  // In forward direction we check if control will for sure reach JoinBB from
+  // InitBB, thus it can not be "stopped" along the way. Ways to "stop" control
+  // are: infinite loops and instructions that do not necessarily transfer
+  // execution to their successor. To check for them we traverse the CFG from
+  // the adjacent blocks to the JoinBB, looking at all intermediate blocks.
+
+  // If we know the function is "will-return" and "no-throw" there is no need
+  // for futher checks.
+  if (!F.hasFnAttribute(Attribute::WillReturn) || !F.doesNotThrow()) {
+
+    auto BlockTransfersExecutionToSuccessor = [](const BasicBlock *BB) {
+      return isGuaranteedToTransferExecutionToSuccessor(BB);
+    };
+
+    SmallPtrSet<const BasicBlock *, 16> Visited;
+    while (!Worklist.empty()) {
+      const BasicBlock *ToBB = Worklist.pop_back_val();
+      if (ToBB == JoinBB)
+        continue;
+
+      // Make sure all loops in-between are finite.
+      if (!Visited.insert(ToBB).second) {
+        if (!F.hasFnAttribute(Attribute::WillReturn)) {
+          bool MayContainIrreducibleControl = getOrCreateCachedOptional(
+              &F, IrreducibleControlMap, mayContainIrreducibleControl, F, LI);
+          if (MayContainIrreducibleControl)
+            return nullptr;
+
+          const Loop *L = LI->getLoopFor(ToBB);
+          if (L && maybeEndlessLoop(*L))
+            return nullptr;
+        }
+
+        continue;
+      }
+
+      // Make sure the block has no instructions that could stop control
+      // transfer.
+      bool TransfersExecution = getOrCreateCachedOptional(
+          ToBB, BlockTransferMap, BlockTransfersExecutionToSuccessor, ToBB);
+      if (!TransfersExecution)
+        return nullptr;
+
+      for (const BasicBlock *AdjacentBB : successors(ToBB))
+        Worklist.push_back(AdjacentBB);
+    }
+  }
+
+  LLVM_DEBUG(dbgs() << "\tJoin block: " << JoinBB->getName() << "\n");
+  return JoinBB;
+}
+
 const Instruction *
 MustBeExecutedContextExplorer::getMustBeExecutedNextInstruction(
     MustBeExecutedIterator &It, const Instruction *PP) {
   if (!PP)
     return PP;
   LLVM_DEBUG(dbgs() << "Find next instruction for " << *PP << "\n");
 
   // If we explore only inside a given basic block we stop at terminators.
MustBeExecutedContextExplorer::getMustBeExecutedNextInstruction(
   // A terminator with a single successor, we will continue at the beginning of
   // that one.
   if (PP->getNumSuccessors() == 1) {
     LLVM_DEBUG(
         dbgs() << "\tUnconditional terminator, continue with successor\n");
     return &PP->getSuccessor(0)->front();
   }
 
+  // Multiple successors mean we need to find the join point where control flow
+  // converges again. We use the findForwardJoinPoint helper function with
+  // information about the function and helper analyses, if available.
+  if (const BasicBlock *JoinBB = findForwardJoinPoint(PP->getParent()))
+    return &JoinBB->front();
+
   LLVM_DEBUG(dbgs() << "\tNo join point found\n");
   return nullptr;
 }
 
 MustBeExecutedIterator::MustBeExecutedIterator(
     MustBeExecutedContextExplorer &Explorer, const Instruction *I)
     : Explorer(Explorer), CurInst(I) {
   reset(I);

llvm/test/Analysis/MustExecute/must_be_executed_context.ll

 define void @simple_conditional(i32 %arg) {
 bb:
   call void @A()
 ; MBEC:      -- Explore context of:   call void @A()
 ; MBEC-NEXT:   [F: simple_conditional]   call void @A()
 ; MBEC-NEXT:   [F: simple_conditional]   call void @B()
 ; MBEC-NEXT:   [F: simple_conditional]   %tmp = icmp eq i32 %arg, 0
 ; MBEC-NEXT:   [F: simple_conditional]   br i1 %tmp, label %bb2, label %bb1
+; MBEC-NEXT:   [F: simple_conditional]   call void @E()
+; MBEC-NEXT:   [F: simple_conditional]   call void @F()
 ; MBEC-NOT:    call
 
   call void @B()
 ; MBEC:      -- Explore context of:   call void @B()
 ; MBEC-NEXT:   [F: simple_conditional]   call void @B()
 ; MBEC-NEXT:   [F: simple_conditional]   %tmp = icmp eq i32 %arg, 0
 ; MBEC-NEXT:   [F: simple_conditional]   br i1 %tmp, label %bb2, label %bb1
+; MBEC-NEXT:   [F: simple_conditional]   call void @E()
+; MBEC-NEXT:   [F: simple_conditional]   call void @F()
 ; MBEC-NOT:    call
 ; MBEC:      -- Explore context of: %tmp
 
   %tmp = icmp eq i32 %arg, 0
   br i1 %tmp, label %bb2, label %bb1
 
 bb1:                                              ; preds = %bb
   call void @C()

llvm/test/Transforms/FunctionAttrs/nonnull.ll

 ; TEST 17 explore child BB test
 ; if(..)
 ;    ... (willreturn & nounwind)
 ; else
 ;    ... (willreturn & nounwind)
 ; fun1(nonnull %a)
 ; We can say that %a is nonnull
 define void @f17(i8* %a, i8 %c) {
-; FIXME: missing nonnull on %a
-; ATTRIBUTOR: define void @f17(i8* %a, i8 %c)
+; ATTRIBUTOR: define void @f17(i8* nonnull %a, i8 %c)
   %cmp = icmp eq i8 %c, 0
   br i1 %cmp, label %if.then, label %if.else
 if.then:
   tail call void @fun0()
   br label %cont
 if.else:
   tail call void @fun0()
   br label %cont
 cont:
   tail call void @fun1(i8* nonnull %a)
   ret void
 }
 ; TEST 18 More complex test
 ; if(..)
 ;    ... (willreturn & nounwind)
 ; else
 ;    ... (willreturn & nounwind)
 ; if(..)
 ;    ... (willreturn & nounwind)
 ; else
 ;    ... (willreturn & nounwind)
 ; fun1(nonnull %a)
 
 define void @f18(i8* %a, i8* %b, i8 %c) {
-; FIXME: missing nonnull on %a
-; ATTRIBUTOR: define void @f18(i8* %a, i8* %b, i8 %c)
+; ATTRIBUTOR: define void @f18(i8* nonnull %a, i8* %b, i8 %c)
   %cmp1 = icmp eq i8 %c, 0
   br i1 %cmp1, label %if.then, label %if.else
 if.then:
   tail call void @fun0()
   br label %cont
 if.else:
   tail call void @fun0()
   br label %cont