Index: lib/Transforms/Scalar/LoopUnswitch.cpp
===================================================================
--- lib/Transforms/Scalar/LoopUnswitch.cpp
+++ lib/Transforms/Scalar/LoopUnswitch.cpp
@@ -212,6 +212,8 @@
     /// Update the appropriate Phi nodes as we do so.
     void SplitExitEdges(Loop *L, const SmallVectorImpl<BasicBlock *> &ExitBlocks);
 
+    bool TryTrivialLoopUnswitch(bool &Changed);
+
     bool UnswitchIfProfitable(Value *LoopCond, Constant *Val,
                               TerminatorInst *TI = nullptr);
     void UnswitchTrivialCondition(Loop *L, Value *Cond, Constant *Val,
@@ -460,6 +462,13 @@
           AC))
     return false;
 
+  // Trivial unswitch condition can only occur at loop header basic block.
+  // Try trivial unswitch first before loop over other basic blocks in the loop.
+  if (TryTrivialLoopUnswitch(Changed)) {
+    ++NumBranches;
+    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.
@@ -668,15 +677,6 @@
 bool LoopUnswitch::UnswitchIfProfitable(Value *LoopCond, Constant *Val,
                                         TerminatorInst *TI) {
   Function *F = loopHeader->getParent();
-  Constant *CondVal = nullptr;
-  BasicBlock *ExitBlock = nullptr;
-
-  if (IsTrivialUnswitchCondition(LoopCond, &CondVal, &ExitBlock)) {
-    // If the condition is trivial, always unswitch. There is no code growth
-    // for this case.
-    UnswitchTrivialCondition(currentLoop, LoopCond, CondVal, ExitBlock, TI);
-    return true;
-  }
 
   // Check to see if it would be profitable to unswitch current loop.
   if (!BranchesInfo.CostAllowsUnswitching()) {
@@ -830,6 +830,45 @@
   ++NumTrivial;
 }
 
+// TryTrivialLoopUnswitch - Check if loop header block's terminator is a trivial
+// unswitch condition (loop header block is the only possible block to have a
+// trivial unswitch condition). If it is, always unswitch becayse there is no
+// code growth for this case.
+bool LoopUnswitch::TryTrivialLoopUnswitch(bool &Changed) {
+  BasicBlock *Header = currentLoop->getHeader();
+  TerminatorInst *HeaderTerm = Header->getTerminator();
+
+  Constant *CondVal = nullptr;
+  BasicBlock *ExitBlock = 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()) {
+      // See if this, or some part of it, is loop invariant.  If so, we can
+      // unswitch on it if we desire.
+      Value *LoopCond = FindLIVLoopCondition(BI->getCondition(),
+                                             currentLoop, Changed);
+      if (LoopCond &&
+          IsTrivialUnswitchCondition(LoopCond, &CondVal, &ExitBlock)) {
+        UnswitchTrivialCondition(currentLoop, LoopCond, CondVal, ExitBlock, HeaderTerm);
+        return true;
+      }
+    }
+  } else if (SwitchInst *SI = dyn_cast<SwitchInst>(HeaderTerm)) {
+    Value *LoopCond = FindLIVLoopCondition(SI->getCondition(),
+                                           currentLoop, Changed);
+    unsigned NumCases = SI->getNumCases();
+    if (LoopCond && NumCases) {
+      if (IsTrivialUnswitchCondition(LoopCond, &CondVal, &ExitBlock)) {
+        UnswitchTrivialCondition(currentLoop, LoopCond, CondVal, ExitBlock, nullptr);
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
 /// SplitExitEdges - Split all of the edges from inside the loop to their exit
 /// blocks.  Update the appropriate Phi nodes as we do so.
 void LoopUnswitch::SplitExitEdges(Loop *L,