[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

770

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.

777

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

803

space missing after the colon.

805

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

770

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.

777

Will do.

803

Will do.

805

Will do.

[• chenli]

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

[broune]

lib/Transforms/Scalar/LoopUnswitch.cpp

752

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.

772

until find -> until it finds

775

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.

783

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

787

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

791

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

752

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

[• chenli]

lib/Transforms/Scalar/LoopUnswitch.cpp

791

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

[broune]

lib/Transforms/Scalar/LoopUnswitch.cpp

791

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

751

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.

771

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

void LoopUnswitch::UnswitchTrivialCondition(Loop *L, Value *Cond, Constant *Val,
   ++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.
 ///
 bool LoopUnswitch::TryTrivialLoopUnswitch(bool &Changed) {
   BasicBlock *Header = currentLoop->getHeader();
   TerminatorInst *HeaderTerm = Header->getTerminator();
   LLVMContext &Context = Header->getContext();
 
   // 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)
     if (I->mayHaveSideEffects())
       return false;
 
+  // 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 should be
+  // fused to loop header block. The reason to not constant folding
+  // conditions 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. This
+  // approach keeps traversing loop header's successors until find the
+  // trivial condition candidate (condition that is not a constant).
+  std::set<BasicBlock*> Visited;
+  Visited.insert(Header);
+
+  // FIXME: add check for constant foldable switch instructions.
+  while (BranchInst *BI = dyn_cast<BranchInst>(HeaderTerm)) {
+    BasicBlock *Successor = nullptr;
+    if (BI->isUnconditional()) {
+      Successor = BI->getSuccessor(0);
+    } else if (BI->getCondition() == ConstantInt::getTrue(Context)) {
+      Successor = BI->getSuccessor(0);
+    } else if (BI->getCondition() == ConstantInt::getFalse(Context)) {
+      Successor = BI->getSuccessor(1);
+    } else {
+      break;
+    }
+
+    assert(Successor && "Does not have a successor?");
+    // 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(Successor) || !Visited.insert(Successor).second)
+      return false;
+
+    // Check if the successor block will execute any side-effecting
+    // instructions.
+    for (BasicBlock::iterator I :*Successor)
+      if (I->mayHaveSideEffects())
+        return false;
+
+    HeaderTerm = Successor->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 this isn't branching on an invariant condition, we can't unswitch it.
     if (!BI->isConditional())

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