[LoopUnswitch] Improve loop unswitch pass to find trivial unswitch conditions more effectively

[• chenli]

Summary

This patch improves trivial loop unswitch.

The current trivial loop unswitch only checks if loop header's terminator contains a trivial unswitch condition. But if the loop header only has one reachable successor (due to intentionally or unintentionally missed code simplification), we should consider the successor as part of the loop header. Therefore, instead of stopping at loop header's terminator, we should keep traversing its successors within loop until reach a *real* conditional branch or switch (whose condition can not be constant folded). This change will enable a single -loop-unswitch pass to unswitch multiple trivial conditions (unswitch one trivial condition could open opportunity to unswitch another one in the same loop), while the old implementation can unswitch only one per pass.

[broune]

lib/Transforms/Scalar/LoopUnswitch.cpp

784

It's not necessarily possible to combine the header with its successor even if it is an unconditional branch, since there might be other predecessors of the successor. It might also be good to point out that unswitching generates branches with constant conditions, so this is expected and not just some corner case.

791

Set can make an allocation per insertion, which is slow. I'd consider SmallSet instead.

http://llvm.org/docs/ProgrammersManual.html#set-like-containers-std-set-smallset-setvector-etc

817

space missing after the colon.

819

This loop is duplicated from outside the loop. Would it make sense to move this loop to the top to to avoid the duplication? It would require a slightly different loop condition. The insert into Visited could also be de-duplicated.

[• chenli]

lib/Transforms/Scalar/LoopUnswitch.cpp

784

Well, I think you can always combine them with some kind of successor duplications if it has multiple predecessors. But it will add too much complexity and it's not worth taking that approach :)

I will add you second comment in code.

791

Will do.

817

Will do.

819

Will do.

[• chenli]

Update patch w.r.t broune's comments.

[broune]

lib/Transforms/Scalar/LoopUnswitch.cpp

749

This comment is correct if we consider the header to extend across several basic blocks, as you do later on in a comment. I suspect that that's not a standard definition of the loop header (correct me if I'm wrong?), making this comment tricky to understand correctly. I think it would be clearer to explain this in terms of the first non-constant branch in the loop. Then probably the "Header" variable should be renamed to something else, like BB. Though this is subjective, so if you prefer the current way, I'll LGTM that.

761

until find -> until it finds

789

This wording suggests that the header is changed in some way, though that's not the case. I'm not sure you need a comment here, but if you keep it, maybe word it as something like "For this function, we consider the header to extend to its successor if ..." and add a full stop at the end.

797

Find -> Found. Add full stop at the end.

801

Nit: This is not necessarily a switch. I think it would be OK without a comment here.

805

Nit: HeaderTerm is set twice and only used twice. I would use Header->getTerminator() directly for the two uses instead, as then there's one variable less to keep track of.

[• chenli]

lib/Transforms/Scalar/LoopUnswitch.cpp

749

Thanks for the suggestion. I should have made the comments consistent with each other. I will revise this :)

[• chenli]

lib/Transforms/Scalar/LoopUnswitch.cpp

805

There are a few uses in the later code as well. I would prefer to keep the variable.

[broune]

lib/Transforms/Scalar/LoopUnswitch.cpp

805

You're right, I thought the function ended before it actually does.

[• chenli]

Update patch w.r.t broune's comments.

[broune]

LGTM.

lib/Transforms/Scalar/LoopUnswitch.cpp

744

Nit: This is correct in the sense of "starting from", but could also be read as "from within". I'd use "starting from" to clarify it.

760

Nit: Now that you are not specifying that this function only looks in the loop header, there is no need to justify looking at successors by saying that they could be seen as part of the loop header.

lib/Transforms/Scalar/LoopUnswitch.cpp

bool LoopUnswitch::processCurrentLoop() {
   // Probably we reach the quota of branches for this loop. If so
   // stop unswitching.
   if (!BranchesInfo.countLoop(
           currentLoop, getAnalysis<TargetTransformInfoWrapperPass>().getTTI(
                            *currentLoop->getHeader()->getParent()),
           AC))
     return false;
 
-  // Trivial unswitch condition can only occur at loop header basic block because
-  // that condition needs to control whether or not the loop does anything at all.
   // Try trivial unswitch first before loop over other basic blocks in the loop.
   if (TryTrivialLoopUnswitch(Changed)) {
     return true;
   }
 
   // Loop over all of the basic blocks in the loop.  If we find an interior
   // block that is branching on a loop-invariant condition, we can unswitch this
   // loop.
void LoopUnswitch::UnswitchTrivialCondition(Loop *L, Value *Cond, Constant *Val,
 
   // Now that we know that the loop is never entered when this condition is a
   // particular value, rewrite the loop with this info.  We know that this will
   // at least eliminate the old branch.
   RewriteLoopBodyWithConditionConstant(L, Cond, Val, false);
   ++NumTrivial;
 }
 
-/// TryTrivialLoopUnswitch - Check if loop header block's terminator is a trivial
-/// unswitch condition: that is, the condition controls whether or not the loop
-/// does anything at all (which means it can only occur in loop header). If it is
-/// a trivial condition, unswitching produces no code duplications (equivalently,
-/// it produces a simpler loop and a new empty loop, which gets deleted). Therefore
-/// always unswitch trivial condition.
+/// TryTrivialLoopUnswitch - Check if the first non-constant condition from the loop
+/// header is a trivial unswitch condition: that is, a condition controls whether
+/// or not the loop does anything at all. If it is a trivial condition, unswitching
+/// produces no code duplications (equivalently, it produces a simpler loop and a new
+/// empty loop, which gets deleted). Therefore always unswitch trivial condition.
 ///
 bool LoopUnswitch::TryTrivialLoopUnswitch(bool &Changed) {
-  BasicBlock *Header = currentLoop->getHeader();
-  TerminatorInst *HeaderTerm = Header->getTerminator();
-  LLVMContext &Context = Header->getContext();
+  BasicBlock *CurrentBB = currentLoop->getHeader();
+  TerminatorInst *CurrentTerm = CurrentBB->getTerminator();
+  LLVMContext &Context = CurrentBB->getContext();
+
+  // If loop header has only one reachable successor (currently via an
+  // unconditional branch or constant foldable conditional branch, but
+  // should also consider adding constant foldable switch instruction in
+  // future), we should keep looking for trivial condition candidates in
+  // the successor as well. Because in theory these successors could be
+  // fused to loop header block. The reason to not constant fold
+  // conditions and merge blocks in LoopUnswitch pass is that it could
+  // potentially break LoopPassManager's invariants. Folding dead branches
+  // could either eliminate the current loop or make other loops unreachable.
+  // LCSSA form might also not be preserved after deleting branches. The
+  // following code keeps traversing loop header's successors until it finds
+  // the trivial condition candidate (condition that is not a constant).
+  // Since unswitching generates branches with constant conditions, this
+  // scenario could be very common in practice.
+  SmallSet<BasicBlock*, 8> Visited;
+
+  while (true) {
+    // If we exit loop or reach a previous visited block, then
+    // we can not reach any trivial condition candidates (unfoldable
+    // branch instructions or switch instructions) and no unswitch
+    // can happen. Exit and return false.
+    if (!currentLoop->contains(CurrentBB) || !Visited.insert(CurrentBB).second)
+      return false;
 
     // Check if this loop will execute any side-effecting instructions (e.g.
     // stores, calls, volatile loads) in the part of the loop that the code
     // *would* execute. Check the header first.
-  for (BasicBlock::iterator I :*Header)
+    for (BasicBlock::iterator I : *CurrentBB)
       if (I->mayHaveSideEffects())
         return false;
 
+    // FIXME: add check for constant foldable switch instructions.
+    if (BranchInst *BI = dyn_cast<BranchInst>(CurrentTerm)) {
+      if (BI->isUnconditional()) {
+        CurrentBB = BI->getSuccessor(0);
+      } else if (BI->getCondition() == ConstantInt::getTrue(Context)) {
+        CurrentBB = BI->getSuccessor(0);
+      } else if (BI->getCondition() == ConstantInt::getFalse(Context)) {
+        CurrentBB = BI->getSuccessor(1);
+      } else {
+        // Found a trivial condition candidate (non-foldable conditional branch).
+        break;
+      }
+    } else {
+      break;
+    }
+
+    CurrentTerm = CurrentBB->getTerminator();
+  }
+
   // CondVal is the condition that controls the trivial condition.
   // LoopExitBB is the BasicBlock that loop exits when meets trivial condition.
   Constant *CondVal = nullptr;
   BasicBlock *LoopExitBB = nullptr;
 
-  if (BranchInst *BI = dyn_cast<BranchInst>(HeaderTerm)) {
+  if (BranchInst *BI = dyn_cast<BranchInst>(CurrentTerm)) {
     // If this isn't branching on an invariant condition, we can't unswitch it.
     if (!BI->isConditional())
       return false;
 
     Value *LoopCond = FindLIVLoopCondition(BI->getCondition(),
                                            currentLoop, Changed);
 
     // Unswitch only if the trivial condition itself is an LIV (not
     // partial LIV which could occur in and/or)
     if (!LoopCond || LoopCond != BI->getCondition())
       return false;
 
     // Check to see if a successor of the branch is guaranteed to
     // exit through a unique exit block without having any
     // side-effects.  If so, determine the value of Cond that causes
     // it to do this.
     if ((LoopExitBB = isTrivialLoopExitBlock(currentLoop,
                                              BI->getSuccessor(0)))) {
       CondVal = ConstantInt::getTrue(Context);
     } else if ((LoopExitBB = isTrivialLoopExitBlock(currentLoop,
                                                     BI->getSuccessor(1)))) {
       CondVal = ConstantInt::getFalse(Context);
     }
 
     // If we didn't find a single unique LoopExit block, or if the loop exit block
     // contains phi nodes, this isn't trivial.
     if (!LoopExitBB || isa<PHINode>(LoopExitBB->begin()))
       return false;   // Can't handle this.
 
-    UnswitchTrivialCondition(currentLoop, LoopCond, CondVal, LoopExitBB, HeaderTerm);
+    UnswitchTrivialCondition(currentLoop, LoopCond, CondVal, LoopExitBB, CurrentTerm);
     ++NumBranches;
     return true;
-  } else if (SwitchInst *SI = dyn_cast<SwitchInst>(HeaderTerm)) {
+  } else if (SwitchInst *SI = dyn_cast<SwitchInst>(CurrentTerm)) {
     // If this isn't switching on an invariant condition, we can't unswitch it.
     Value *LoopCond = FindLIVLoopCondition(SI->getCondition(),
                                            currentLoop, Changed);
 
     // Unswitch only if the trivial condition itself is an LIV (not
     // partial LIV which could occur in and/or)
     if (!LoopCond || LoopCond != SI->getCondition())
       return false;

test/Transforms/LoopUnswitch/infinite-loop.ll

 ; REQUIRES: asserts
 ; RUN: opt -loop-unswitch -disable-output -stats -info-output-file - < %s | FileCheck --check-prefix=STATS %s
 ; RUN: opt -loop-unswitch -simplifycfg -S < %s | FileCheck %s
 ; PR5373
 
 ; Loop unswitching shouldn't trivially unswitch the true case of condition %a
 ; in the code here because it leads to an infinite loop. While this doesn't
 ; contain any instructions with side effects, it's still a kind of side effect.
 ; It can trivially unswitch on the false cas of condition %a though.
 
 ; STATS: 2 loop-unswitch - Number of branches unswitched
-; STATS: 1 loop-unswitch - Number of unswitches that are trivial
+; STATS: 2 loop-unswitch - Number of unswitches that are trivial
 
 ; CHECK-LABEL: @func_16(
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT: br i1 %a, label %entry.split, label %abort0.split
 
 ; CHECK: entry.split:
-; CHECK-NEXT: br i1 %b, label %cond.end.us, label %abort1
+; CHECK-NEXT: br i1 %b, label %cond.end, label %abort1.split
 
-; CHECK: cond.end.us:
-; CHECK-NEXT: br label %cond.end.us
+; CHECK: cond.end:
+; CHECK-NEXT: br label %cond.end
 
 ; CHECK: abort0.split:
 ; CHECK-NEXT: call void @end0() [[NOR_NUW:#[0-9]+]]
 ; CHECK-NEXT: unreachable
 
-; CHECK: abort1:
+; CHECK: abort1.split:
 ; CHECK-NEXT: call void @end1() [[NOR_NUW]]
 ; CHECK-NEXT: unreachable
 
 ; CHECK: }
 
 define void @func_16(i1 %a, i1 %b) nounwind {
 entry:
   br label %for.body

test/Transforms/LoopUnswitch/trivial-unswitch.ll

+; RUN: opt < %s -loop-unswitch -loop-unswitch-threshold=0 -verify-loop-info -S < %s 2>&1 | FileCheck %s
+
+; This test contains two trivial unswitch condition in one loop. 
+; LoopUnswitch pass should be able to unswitch the second one 
+; after unswitching the first one.
+
+
+; CHECK:  br i1 %cond1, label %..split_crit_edge, label %.loop_exit.split_crit_edge
+
+; CHECK:  ..split_crit_edge:                                ; preds = %0
+; CHECK:    br label %.split
+
+; CHECK:  .split:                                           ; preds = %..split_crit_edge
+; CHECK:    br i1 %cond2, label %.split..split.split_crit_edge, label %.split.loop_exit.split1_crit_edge
+
+; CHECK:  .split..split.split_crit_edge:                    ; preds = %.split
+; CHECK:    br label %.split.split
+
+; CHECK:  .split.split:                                     ; preds = %.split..split.split_crit_edge
+; CHECK:    br label %loop_begin
+
+; CHECK:  loop_begin:                                       ; preds = %do_something, %.split.split
+; CHECK:    br i1 true, label %continue, label %loop_exit
+
+; CHECK:  continue:                                         ; preds = %loop_begin
+; CHECK:    %var_val = load i32, i32* %var
+; CHECK:    br i1 true, label %do_something, label %loop_exit
+
+define i32 @test(i32* %var, i1 %cond1, i1 %cond2) {
+  br label %loop_begin
+
+loop_begin:  
+  br i1 %cond1, label %continue, label %loop_exit  ; first trivial condition
+
+continue:
+  %var_val = load i32, i32* %var
+  br i1 %cond2, label %do_something, label %loop_exit  ; second trivial condition  
+
+do_something:
+  call void @some_func() noreturn nounwind
+  br label %loop_begin
+
+loop_exit:
+  ret i32 0
+}
+
+declare void @some_func() noreturn
 No newline at end of file