[LoopDeletion] Handle switch in proving that loop exits on first iteration

[dmakogon]

Summary

Added check for switch-terminated blocks in loops.
Now if a block is terminated with a switch, we try to find out which of the cases is taken on 1st iteration and mark corresponding edge from the block to the case successor as live.

[mkazantsev]

Tests?

[dmakogon]

Updated test checks

[mkazantsev]

Fine by me, but please someone else take a look.

[nikic]

LGTM as well.

llvm/lib/Transforms/Scalar/LoopDeletion.cpp

321–326

Just as a side-note: I think it would be nice to move the SimplifyICmpInst handling into getValueOnFirstIteration(), and then this could just directly work on the br argument.

[mkazantsev]

Good idea, we'll follow-up with this as NFC.

llvm/lib/Transforms/Scalar/LoopDeletion.cpp

for (auto &PN : BB->phis()) {
       }
     }
 
     using namespace PatternMatch;
     ICmpInst::Predicate Pred;
     Value *LHS, *RHS;
     BasicBlock *IfTrue, *IfFalse;
     auto *Term = BB->getTerminator();
-    // TODO: Handle switch.
-    if (!match(Term, m_Br(m_ICmp(Pred, m_Value(LHS), m_Value(RHS)),
-                          m_BasicBlock(IfTrue), m_BasicBlock(IfFalse)))) {
-      MarkAllSuccessorsLive(BB);
-      continue;
-    }
-
+    if (match(Term, m_Br(m_ICmp(Pred, m_Value(LHS), m_Value(RHS)),
       if (!LHS->getType()->isIntegerTy()) {
         MarkAllSuccessorsLive(BB);
         continue;
       }
 
       // Can we prove constant true or false for this condition?
       LHS = getValueOnFirstIteration(LHS, FirstIterValue, SQ);
       RHS = getValueOnFirstIteration(RHS, FirstIterValue, SQ);
       auto *KnownCondition = SimplifyICmpInst(Pred, LHS, RHS, SQ);
       if (!KnownCondition) {
         // Failed to simplify.
         MarkAllSuccessorsLive(BB);
         continue;
       }
       if (isa<UndefValue>(KnownCondition)) {
         // TODO: According to langref, branching by undef is undefined behavior.
         // It means that, theoretically, we should be able to just continue
         // without marking any successors as live. However, we are not certain
         // how correct our compiler is at handling such cases. So we are being
         // very conservative here.
         //
         // If there is a non-loop successor, always assume this branch leaves the
         // loop. Otherwise, arbitrarily take IfTrue.
         //
         // Once we are certain that branching by undef is handled correctly by
         // other transforms, we should not mark any successors live here.
         if (L->contains(IfTrue) && L->contains(IfFalse))
           MarkLiveEdge(BB, IfTrue);
         continue;
       }
       auto *ConstCondition = dyn_cast<ConstantInt>(KnownCondition);
       if (!ConstCondition) {
         // Non-constant condition, cannot analyze any further.
         MarkAllSuccessorsLive(BB);
         continue;
       }
       if (ConstCondition->isAllOnesValue())
         MarkLiveEdge(BB, IfTrue);
       else
         MarkLiveEdge(BB, IfFalse);
+    } else if (SwitchInst *SI = dyn_cast<SwitchInst>(Term)) {
+      auto *SwitchValue = SI->getCondition();
+      auto *SwitchValueOnFirstIter =
+          getValueOnFirstIteration(SwitchValue, FirstIterValue, SQ);
+      auto *ConstSwitchValue = dyn_cast<ConstantInt>(SwitchValueOnFirstIter);
+      if (!ConstSwitchValue) {
+        MarkAllSuccessorsLive(BB);
+        continue;
+      }
+      auto CaseIterator = SI->findCaseValue(ConstSwitchValue);
+      MarkLiveEdge(BB, CaseIterator->getCaseSuccessor());
+    } else {
+      MarkAllSuccessorsLive(BB);
+      continue;
+    }
   }
 
   // We can break the latch if it wasn't live.
   return !LiveEdges.count({ Latch, Header });
 }
 
 /// If we can prove the backedge is untaken, remove it.  This destroys the
 /// loop, but leaves the (now trivially loop invariant) control flow and

llvm/test/Transforms/LoopDeletion/eval_first_iteration.ll

 ; Switch tests
 
 ; Here switch will always jump to the default label
 define i32 @test_switch_ne_default() {
 ; CHECK-LABEL: @test_switch_ne_default(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[SUM_NEXT:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[SUB:%.*]] = sub i32 4, [[SUM]]
 ; CHECK-NEXT:    switch i32 [[SUB]], label [[DEFAULT:%.*]] [
 ; CHECK-NEXT:    i32 0, label [[ONZERO:%.*]]
 ; CHECK-NEXT:    i32 1, label [[ONONE:%.*]]
 ; CHECK-NEXT:    i32 2, label [[ONTWO:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       default:
-; CHECK-NEXT:    br label [[BACKEDGE]]
+; CHECK-NEXT:    br label [[BACKEDGE:%.*]]
 ; CHECK:       onzero:
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       onone:
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       ontwo:
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[MERGE_PHI:%.*]] = phi i32 [ [[SUB]], [[DEFAULT]] ], [ 0, [[ONZERO]] ], [ 1, [[ONONE]] ], [ 2, [[ONTWO]] ]
-; CHECK-NEXT:    [[SUM_NEXT]] = add i32 [[SUM]], [[MERGE_PHI]]
+; CHECK-NEXT:    [[SUM_NEXT:%.*]] = add i32 [[SUM]], [[MERGE_PHI]]
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp ne i32 [[SUM_NEXT]], 4
-; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[DONE:%.*]]
+; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[BACKEDGE_LOOP_CRIT_EDGE:%.*]], label [[DONE:%.*]]
+; CHECK:       backedge.loop_crit_edge:
+; CHECK-NEXT:    unreachable
 ; CHECK:       done:
 ; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT]], [[BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
 ;
 entry:
   br label %loop
 
 loop:                                             ; preds = %backedge, %entry

 }
 
 ; Here switch will always jump to the %ontwo label
 define i32 @test_switch_ne_one_case() {
 ; CHECK-LABEL: @test_switch_ne_one_case(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[SUM_NEXT:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[SUB:%.*]] = sub i32 4, [[SUM]]
 ; CHECK-NEXT:    switch i32 [[SUB]], label [[DEFAULT:%.*]] [
 ; CHECK-NEXT:    i32 0, label [[ONZERO:%.*]]
 ; CHECK-NEXT:    i32 1, label [[ONONE:%.*]]
 ; CHECK-NEXT:    i32 4, label [[ONTWO:%.*]]
 ; CHECK-NEXT:    ]
 ; CHECK:       default:
-; CHECK-NEXT:    br label [[BACKEDGE]]
+; CHECK-NEXT:    br label [[BACKEDGE:%.*]]
 ; CHECK:       onzero:
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       onone:
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       ontwo:
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[MERGE_PHI:%.*]] = phi i32 [ 2, [[DEFAULT]] ], [ 0, [[ONZERO]] ], [ 1, [[ONONE]] ], [ [[SUB]], [[ONTWO]] ]
-; CHECK-NEXT:    [[SUM_NEXT]] = add i32 [[SUM]], [[MERGE_PHI]]
+; CHECK-NEXT:    [[SUM_NEXT:%.*]] = add i32 [[SUM]], [[MERGE_PHI]]
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp ne i32 [[SUM_NEXT]], 4
-; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[DONE:%.*]]
+; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[BACKEDGE_LOOP_CRIT_EDGE:%.*]], label [[DONE:%.*]]
+; CHECK:       backedge.loop_crit_edge:
+; CHECK-NEXT:    unreachable
 ; CHECK:       done:
 ; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT]], [[BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
 ;
 entry:
   br label %loop
 
 loop:                                             ; preds = %backedge, %entry

 }
 
 ; Here switch will always jump to the %backedge label, but there are two jumps to this label in switch
 define i32 @test_switch_ne_one_case_identical_jumps() {
 ; CHECK-LABEL: @test_switch_ne_one_case_identical_jumps(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[SUM_NEXT:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[SUM:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[SUB:%.*]] = sub i32 2, [[SUM]]
 ; CHECK-NEXT:    switch i32 [[SUB]], label [[DEFAULT:%.*]] [
 ; CHECK-NEXT:    i32 0, label [[FIRST_BLOCK:%.*]]
-; CHECK-NEXT:    i32 1, label [[BACKEDGE]]
+; CHECK-NEXT:    i32 1, label [[BACKEDGE:%.*]]
 ; CHECK-NEXT:    i32 2, label [[BACKEDGE]]
 ; CHECK-NEXT:    ]
 ; CHECK:       default:
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       first_block:
 ; CHECK-NEXT:    br label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[MERGE_PHI:%.*]] = phi i32 [ 0, [[DEFAULT]] ], [ 1, [[FIRST_BLOCK]] ], [ [[SUB]], [[LOOP]] ], [ [[SUB]], [[LOOP]] ]
-; CHECK-NEXT:    [[SUM_NEXT]] = add i32 [[SUM]], [[MERGE_PHI]]
+; CHECK-NEXT:    [[SUM_NEXT:%.*]] = add i32 [[SUM]], [[MERGE_PHI]]
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp ne i32 [[SUM_NEXT]], 2
-; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[DONE:%.*]]
+; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[BACKEDGE_LOOP_CRIT_EDGE:%.*]], label [[DONE:%.*]]
+; CHECK:       backedge.loop_crit_edge:
+; CHECK-NEXT:    unreachable
 ; CHECK:       done:
 ; CHECK-NEXT:    [[SUM_NEXT_LCSSA:%.*]] = phi i32 [ [[SUM_NEXT]], [[BACKEDGE]] ]
 ; CHECK-NEXT:    ret i32 [[SUM_NEXT_LCSSA]]
 ;
 entry:
   br label %loop
 
 loop:                                             ; preds = %backedge, %entry