Fix a bug when unswitching on partial LIV for SwitchInst

[trentxintong]

Summary

Fix a bug when unswitching on partial LIV for SwitchInst.

[trentxintong]

Reorganize some logic.

[trentxintong]

useless-unswitches-exhaust-quota.pdf32 KBDownload

This is the CFG I get by running unswitching on the 1st test case included in this patch (without the patch), we are unswitching on the value 11 a few times till we've taken up all the quota. (Also learn how to attach files in Phabricator).

[david2050]

Hi Xin, let me make sure I understand given:

for(...;...;++i)
   switch (f(i) & c) 
        case 11:
        case 12:

we might unswitch this too

if (c) 
   for(...;...;++i)
     switch (f(i) & c) 
         case 11:
         case 12:
else
  for(...;...;++i)
   switch (0) 
        case 11:
        case 12:

and we achieve no simplification in one loop.

How does restricting to a single bit change this?

[trentxintong]

In D29107#658375, @david2050 wrote:

Hi Xin, let me make sure I understand given:

for(...;...;++i)
   switch (f(i) & c) 
        case 11:
        case 12:

we might unswitch this too

if (c) 
   for(...;...;++i)
     switch (f(i) & c) 
         case 11:
         case 12:
else
  for(...;...;++i)
   switch (0) 
        case 11:
        case 12:

and we achieve no simplification in one loop.

How does restricting to a single bit change this?

Hi David

We are trying to pick a value for c such that f(i) & c can be simplified to either a loop _variant_ value which we will not try to unswitch on again, or a constant which we will not unswitch either. The problem is when c is more than 1 bit, we can not pick a value that will guarantee such simplifications. Lets say c is a signed int here and we pick a random value, then we do not know what f(i) & c will simplify to, maybe it will not simplify at all. In such cases, given the way loop unswitching is implemented, it will try to unswitch this loop again (It records what case value has been unswitched and will not unswitch on that case-value again).

In your case, we picked c==0, we get simplification in the else-loop and no simplification in the if-loop as what you have in the example. However, the way loop unswitching is implemented is that we try to keep unswitching the current loop (the if-loop in this case) until it cant be unswitched anymore or no quota left. We know c is not 0, but we do not know what c is and thus f(i) & c can not be simplified, so we do not know whats unswitched out ...

If c is one bit, it can only be true or false, we are guaranteed to either reach a constant or a loop _variant_ input value for the switch for the current loop and the cloned loop.

[trentxintong]

In D29107#658479, @trentxintong wrote:

In D29107#658375, @david2050 wrote:

Hi Xin, let me make sure I understand given:

for(...;...;++i)
   switch (f(i) & c) 
        case 11:
        case 12:

we might unswitch this too

if (c) 
   for(...;...;++i)
     switch (f(i) & c) 
         case 11:
         case 12:
else
  for(...;...;++i)
   switch (0) 
        case 11:
        case 12:

and we achieve no simplification in one loop.

How does restricting to a single bit change this?

Hi David

We are trying to pick a value for c such that f(i) & c can be simplified to either a loop _variant_ value which we will not try to unswitch on again, or a constant which we will not unswitch either. The problem is when c is more than 1 bit, we can not pick a value that will guarantee such simplifications. Lets say c is a signed int here and we pick a random value, then we do not know what f(i) & c will simplify to, maybe it will not simplify at all. In such cases, given the way loop unswitching is implemented, it will try to unswitch this loop again (It records what case value has been unswitched and will not unswitch on that case-value again).

In your case, we picked c==0, we get simplification in the else-loop and no simplification in the if-loop as what you have in the example. However, the way loop unswitching is implemented is that we try to keep unswitching the current loop (the if-loop in this case) until it cant be unswitched anymore or no quota left. We know c is not 0, but we do not know what c is and thus f(i) & c can not be simplified, so we do not know whats unswitched out ...

If c is one bit, it can only be true or false, we are guaranteed to either reach a constant or a loop _variant_ input value for the switch for the current loop and the cloned loop.

To be more precise, its not we are picking c == 0, its we are trying to unswitch the following input value f(i) & c (0 in this case) out of the loop.

[trentxintong]

@sanjoy @efriedma @hfinkel Mind taking a look ? Especially this fixes a bug which make loop unswitching less effective. Thanks !.

[sanjoy]

I would say the problem here is that we're not doing a good enough job of choosing a good value to unswitch on (and not that we're unswitching on the wrong "type" of value). That is, say we have a switch on Inv & Var (Inv being the loop invariant value) that dispatches on 0, 200 and 300. In such a case, we should clearly unswitch on Inv == 0 to simplify the switch to a direct branch to the (Inv & Var) == 0 case in one of the loops. Similar logic follows if we have a switch on Inv | Var and one of the cases is -1.

So I think the patch should be structured as:

• Keep FindLIVLoopCondition the same (maybe return a bit of extra information to make the next steps easier)
• In processCurrentLoop, check to see if the expression we're switching on simplifies to any of the case values given some constant value for what was returned by FindLIVLoopCondition
  • If there is such a constant value, unswitch on that and profit
  • Else give up

lib/Transforms/Scalar/LoopUnswitch.cpp

442

Did you clang-format this bit?

[trentxintong]

In D29107#664178, @sanjoy wrote:

I would say the problem here is that we're not doing a good enough job of choosing a good value to unswitch on (and not that we're unswitching on the wrong "type" of value). That is, say we have a switch on Inv & Var (Inv being the loop invariant value) that dispatches on 0, 200 and 300. In such a case, we should clearly unswitch on Inv == 0 to simplify the switch to a direct branch to the (Inv & Var) == 0 case in one of the loops. Similar logic follows if we have a switch on Inv | Var and one of the cases is -1.

So I think the patch should be structured as:

• Keep FindLIVLoopCondition the same (maybe return a bit of extra information to make the next steps easier)
• In processCurrentLoop, check to see if the expression we're switching on simplifies to any of the case values given some constant value for what was returned by FindLIVLoopCondition
  • If there is such a constant value, unswitch on that and profit
  • Else give up

Yes, the ability to simplify f(x) (The chain of AND(s) and OR(s)) to a meaningful case value is important here. Our current way is to pick 1 of the case values and expecting f(x) simplifies to a meaningful case value, which probably does not happen in many cases. Since we are concerning with a chain of and/or operators working on one or multiple loop-variant values and a single loop invariant values, I think this is what we could do, we record wether the chain contains AND and OR. Once we have that we can do the following:

• If we have a case 0 in the switch and we only have AND, we can pick loop_invariant as 0, this way we can unswitch case 0 out of the loop.
• if we have a case UNSIGNED_TYPE_MAX (signed -1) and we only have ORs, we can pick the loop_invariant as UNSIGNED_TYPE_MAX and unswitch UNSIGNED_TYPE_MAX out of the loop.

What do you think ?

[sanjoy]

In D29107#664209, @trentxintong wrote:

Yes, the ability to simplify f(x) (The chain of AND(s) and OR(s)) to a meaningful case value is important here. Our current way is to pick 1 of the case values and expecting f(x) simplifies to a meaningful case value, which probably does not happen in many cases. Since we are concerning with a chain of and/or operators working on one or multiple loop-variant values and a single loop invariant values, I think this is what we could do, we record wether the chain contains AND and OR. Once we have that we can do the following:

• If we have a case 0 in the switch and we only have AND, we can pick loop_invariant as 0, this way we can unswitch case 0 out of the loop.
• if we have a case UNSIGNED_TYPE_MAX (signed -1) and we only have ORs, we can pick the loop_invariant as UNSIGNED_TYPE_MAX and unswitch UNSIGNED_TYPE_MAX out of the loop.

What do you think ?

Sounds good!

[trentxintong]

Address @sanjoy comments.

[sanjoy]

Comments inline.

lib/Transforms/Scalar/LoopUnswitch.cpp

372

Why not drop the Ty and call it a OperatorChain? I've only seen Ty used in templated code when it isn't 100% obvious that something is a type.

439

Please put braces around such a large if block to make its extent obvious.

469

Doesn't this regress the callers of FindLIVLoopCondition that pass in nullptr for OCS? They used to get the recursive search before, and now they don't?

My preference would be to avoid allowing a null OCS, and instead have it be a OperatorChainTy & instead. Callers that don't need it can pass in a dummy.

681

auto *

[trentxintong]

lib/Transforms/Scalar/LoopUnswitch.cpp

469

The only places which we pass in a nullptr is TryTrivialLoopUnswitch (with the patch), i.e. before the patch TryTrivialLoopUnswitch ignores partial LIV even we manage to get one by walking up the chain. So there is no place which will get regressed in the current state the pass is written.

[trentxintong]

Address comments from @sanjoy. There is still one unresolved issue, waiting to be discussed and addressed.

[trentxintong]

I took @sanjoy suggestion about passing in the OperatorChain as a reference, instead of pointer (and allowing nullptr to be passed) in.

[trentxintong]

Gentle ping.

[sanjoy]

Mostly minor comments.

lib/Transforms/Scalar/LoopUnswitch.cpp

373

Normally I've seen LLVM prefixes these enums as OC_OpChainNone etc.

420

I'd leave this uninitialized here. That way if you forget to assign to it down one path, sanitizers / compiler warnings will catch the bug.

I'm also not sold on the OCS and CChain names -- how about ParentChain and NewChain?

431

I'd rather have this last case be:

case OpChainMixed:
  CChain = OpChainMixed;
  break;

That will make this switch "fully covered" and have the compiler issue a warning if someone adds a new enum but forgets to update the switch.

456

This style is a bit awkward, IMO. (In a later change, no need to pre-commit review) do you mind changing FindLIVLoopCondition to return std::pair<Value *, OperatorChain> as well?

683

Why do you care about the case being selected not being the default?

test/Transforms/LoopUnswitch/basictest.ll

186

s/dont/don't/
s/a input/an input/

[trentxintong]

lib/Transforms/Scalar/LoopUnswitch.cpp

683

In its form, loop unswitch does not unswitch any value out of the default case. If we unswitch something out of the default here, it will not be recorded as unswitched-out and it will try to unswitch it again.

I guess this can be improved, but i prefer in another patch.

[trentxintong]

Address @sanjoy comments.

lib/Transforms/Scalar/LoopUnswitch.cpp

456

Yes, I will do this in a subsequent commit.

[trentxintong]

Decided to teach how to unswitch value out of default-case in case its can lead to
simplifications in at least one of the loops.

Also fix a small bug about how unswitched value is recorded. Updated the test cases to
check for that.

[sanjoy]

LGTM, but let's not fold in the changes to RewriteLoopBodyWithConditionConstant into this patch.

lib/Transforms/Scalar/LoopUnswitch.cpp

684

Not sure what you meany by !.. If it is supposed to be sentence terminator then I'd just go with a period. Same below.

1338

This bit is new right? Let's do this change in a different commit and go with the bailout on a default switch case in FindLIVLoopCondition in the first step (but please add a FIXME to FindLIVLoopCondition).

[trentxintong]

lib/Transforms/Scalar/LoopUnswitch.cpp

1338

Actually this change is necessary to make the patch work. When I first wrote this patch to solve the problem, I was disabling walking up the operator chain for switches. And this would work and I manually checked the test case worked.

Then we changed the patch to move up the chain and now the problem is: If we have a chain of ands and we have a partial LIV on top of the chain, we will not get to BranchesInfo.setUnswitched(SI, Val); here. Because the SwitchInst is not a direct user of the partial LIV.

What i did is to move this setUnswitched out to where SwitchInst is unswitched, i.e. in TrivialUnswitching and NonTrivialUnswitching. And with that we can record what is unswitched. I also added more checks into the test cases to make sure unswitching only happen once.

[trentxintong]

Minor comment fixes.

[trentxintong]

Basically, there are 2 problems here this patch addresses.

1. we try to find a value that actually simplify the operator chain and in case we cant, we give up and this is better than picking a random value and hope it simplifies.
2. we fix how unswitch value is recorded. current way (before the patch) can not record partial LIV.

Both are needed to have the patch working, as we can not allow partial LIV by walking up the chain and not fixing how to record unswitched partial LIV correctly.

[trentxintong]

@sanjoy can you please take a look at this again ? Thanks.

[sanjoy]

lgtm again

lib/Transforms/Scalar/LoopUnswitch.cpp

1060

I understand why you need this, but IIUC the CurLoopInstructions table is unnecessary.

We can solve the problem it solves in either of two ways:

• Remove cases from switches after unswitching them. That is, with switch (V) { case 0: ... } the V != 0 version of the loop does not have a case 0:. This would prevent us from unswitching on that condition again by definition.
• Do not unswitch on cases that lead to a block that solely contains unreachable. This way we will avoid doing redundant work generally (simplifycfg will clean these trivial cases up), and also prevent us unswitching switch es that we just unswitched.

So while your change itself is okay, depending on how much time you have it would be nice to remove CurLoopInstructions altogether.

[trentxintong]

lib/Transforms/Scalar/LoopUnswitch.cpp

1060

• Remove cases from switches after unswitching them. I think this is possible.

• Do not unswitch on cases that lead to a block that solely contains unreachable. - We can not split critical edges with some obscure cases correctly right now, say we have 2 edges from BB-A to BB-B, when we call SplitEdge, we do not know which one is split. To be more specific, the one returned by GetSuccessorNumber(BB, Succ) is split. In such case we can not mark a case with unreachable because we may not get a basic block to put that unreachable instruction. This can be fixed as well.

I am going to land this for now. And think about what i can do to make this better sometime next month or so.

[trentxintong]

@sanjoy Thanks alot for spending time with me to drive this forward.

lib/Transforms/Scalar/LoopUnswitch.cpp

 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(LoopUnswitch, "loop-unswitch", "Unswitch loops",
                       false, false)
 
 Pass *llvm::createLoopUnswitchPass(bool Os) {
   return new LoopUnswitch(Os);
 }
 
+/// Operator chain lattice.
+enum OperatorChain {
+  OC_OpChainNone,    ///< There is no operator.
+  OC_OpChainOr,      ///< There are only ORs.
+  OC_OpChainAnd,     ///< There are only ANDs.
+  OC_OpChainMixed    ///< There are ANDs and ORs.
+};
+
 /// Cond is a condition that occurs in L. If it is invariant in the loop, or has
 /// an invariant piece, return the invariant. Otherwise, return null.
+//
+/// NOTE: FindLIVLoopCondition will not return a partial LIV by walking up a
+/// mixed operator chain, as we can not reliably find a value which will simplify
+/// the operator chain. If the chain is AND-only or OR-only, we can use 0 or ~0
+/// to simplify the chain.
+///
+/// NOTE: In case a partial LIV and a mixed operator chain, we may be able to
+/// simplify the condition itself to a loop variant condition, but at the
+/// cost of creating an entirely new loop.
 static Value *FindLIVLoopCondition(Value *Cond, Loop *L, bool &Changed,
+                                   OperatorChain &ParentChain,
                                    DenseMap<Value *, Value *> &Cache) {
   auto CacheIt = Cache.find(Cond);
   if (CacheIt != Cache.end())
     return CacheIt->second;
 
   // We started analyze new instruction, increment scanned instructions counter.
   ++TotalInsts;
 
   // We can never unswitch on vector conditions.
   if (Cond->getType()->isVectorTy())
     return nullptr;
 
   // Constants should be folded, not unswitched on!
   if (isa<Constant>(Cond)) return nullptr;
 
   // TODO: Handle: br (VARIANT|INVARIANT).
 
   // Hoist simple values out.
   if (L->makeLoopInvariant(Cond, Changed)) {
     Cache[Cond] = Cond;
     return Cond;
   }
 
+  // Walk up the operator chain to find partial invariant conditions.
   if (BinaryOperator *BO = dyn_cast<BinaryOperator>(Cond))
     if (BO->getOpcode() == Instruction::And ||
         BO->getOpcode() == Instruction::Or) {
+      // Given the previous operator, compute the current operator chain status.
+      OperatorChain NewChain;
+      switch (ParentChain) {
+      case OC_OpChainNone:
+        NewChain = BO->getOpcode() == Instruction::And ? OC_OpChainAnd :
+                                      OC_OpChainOr;
+        break;
+      case OC_OpChainOr:
+        NewChain = BO->getOpcode() == Instruction::Or ? OC_OpChainOr :
+                                      OC_OpChainMixed;
+        break;
+      case OC_OpChainAnd:
+        NewChain = BO->getOpcode() == Instruction::And ? OC_OpChainAnd :
+                                      OC_OpChainMixed;
+        break;
+      case OC_OpChainMixed:
+        NewChain = OC_OpChainMixed;
+        break;
+      }
+
+      // If we reach a Mixed state, we do not want to keep walking up as we can not
+      // reliably find a value that will simplify the chain. With this check, we
+      // will return null on the first sight of mixed chain and the caller will
+      // either backtrack to find partial LIV in other operand or return null.
+      if (NewChain != OC_OpChainMixed) {
+        // Update the current operator chain type before we search up the chain.
+        ParentChain = NewChain;
         // If either the left or right side is invariant, we can unswitch on this,
         // which will cause the branch to go away in one loop and the condition to
         // simplify in the other one.
-      if (Value *LHS =
-              FindLIVLoopCondition(BO->getOperand(0), L, Changed, Cache)) {
+        if (Value *LHS = FindLIVLoopCondition(BO->getOperand(0), L, Changed,
+                                              ParentChain, Cache)) {
           Cache[Cond] = LHS;
           return LHS;
         }
-      if (Value *RHS =
-              FindLIVLoopCondition(BO->getOperand(1), L, Changed, Cache)) {
+        // We did not manage to find a partial LIV in operand(0). Backtrack and try
+        // operand(1).
+        ParentChain = NewChain;
+        if (Value *RHS = FindLIVLoopCondition(BO->getOperand(1), L, Changed,
+                                              ParentChain, Cache)) {
           Cache[Cond] = RHS;
           return RHS;
         }
       }
+    }
 
   Cache[Cond] = nullptr;
   return nullptr;
 }
 
-static Value *FindLIVLoopCondition(Value *Cond, Loop *L, bool &Changed) {
+/// Cond is a condition that occurs in L. If it is invariant in the loop, or has
+/// an invariant piece, return the invariant along with the operator chain type.
+/// Otherwise, return null.
+static std::pair<Value *, OperatorChain> FindLIVLoopCondition(Value *Cond,
+                                                              Loop *L,
+                                                              bool &Changed) {
   DenseMap<Value *, Value *> Cache;
-  return FindLIVLoopCondition(Cond, L, Changed, Cache);
+  OperatorChain OpChain = OC_OpChainNone;
+  Value *FCond = FindLIVLoopCondition(Cond, L, Changed, OpChain, Cache);
+
+  // In case we do find a LIV, it can not be obtained by walking up a mixed
+  // operator chain.
+  assert((!FCond || OpChain != OC_OpChainMixed) &&
+        "Do not expect a partial LIV with mixed operator chain");
+  return {FCond, OpChain};
 }
 
 bool LoopUnswitch::runOnLoop(Loop *L, LPPassManager &LPM_Ref) {
   if (skipLoop(L))
     return false;
 
   AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(
       *L->getHeader()->getParent());
bool LoopUnswitch::processCurrentLoop() {
   // Do not do non-trivial unswitch while optimizing for size.
   // FIXME: Use Function::optForSize().
   if (OptimizeForSize ||
       loopHeader->getParent()->hasFnAttribute(Attribute::OptimizeForSize))
     return false;
 
   for (IntrinsicInst *Guard : Guards) {
     Value *LoopCond =
-        FindLIVLoopCondition(Guard->getOperand(0), currentLoop, Changed);
+        FindLIVLoopCondition(Guard->getOperand(0), currentLoop, Changed).first;
     if (LoopCond &&
         UnswitchIfProfitable(LoopCond, ConstantInt::getTrue(Context))) {
       // NB! Unswitching (if successful) could have erased some of the
       // instructions in Guards leaving dangling pointers there.  This is fine
       // because we're returning now, and won't look at Guards again.
       ++NumGuards;
       return true;
     }
if (BranchInst *BI = dyn_cast<BranchInst>(TI)) {
         continue;
  
       // 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);
+                                               currentLoop, Changed).first;
         if (LoopCond &&
             UnswitchIfProfitable(LoopCond, ConstantInt::getTrue(Context), TI)) {
           ++NumBranches;
           return true;
         }
       }
     } else if (SwitchInst *SI = dyn_cast<SwitchInst>(TI)) {
-      Value *LoopCond = FindLIVLoopCondition(SI->getCondition(),
-                                             currentLoop, Changed);
+      Value *SC = SI->getCondition();
+      Value *LoopCond;
+      OperatorChain OpChain;
+      std::tie(LoopCond, OpChain) =
+        FindLIVLoopCondition(SC, currentLoop, Changed);
+
       unsigned NumCases = SI->getNumCases();
       if (LoopCond && NumCases) {
         // Find a value to unswitch on:
         // FIXME: this should chose the most expensive case!
         // FIXME: scan for a case with a non-critical edge?
         Constant *UnswitchVal = nullptr;
-
+        // Find a case value such that at least one case value is unswitched
+        // out.
+        if (OpChain == OC_OpChainAnd) {
+          // If the chain only has ANDs and the switch has a case value of 0.
+          // Dropping in a 0 to the chain will unswitch out the 0-casevalue.
+          auto *AllZero = cast<ConstantInt>(Constant::getNullValue(SC->getType()));
+          if (BranchesInfo.isUnswitched(SI, AllZero))
+            continue;
+          // We are unswitching 0 out !.
+          UnswitchVal = AllZero;
+        } else if (OpChain == OC_OpChainOr) {
+          // If the chain only has ORs and the switch has a case value of ~0.
+          // Dropping in a ~0 to the chain will unswitch out the ~0-casevalue.
+          auto *AllOne = cast<ConstantInt>(Constant::getAllOnesValue(SC->getType()));
+          if (BranchesInfo.isUnswitched(SI, AllOne))
+            continue;
+          // We are unswitching ~0 out !.
+          UnswitchVal = AllOne;
+        } else {
+          assert(OpChain == OC_OpChainNone && 
+                 "Expect to unswitch on trivial chain");
           // Do not process same value again and again.
           // At this point we have some cases already unswitched and
           // some not yet unswitched. Let's find the first not yet unswitched one.
           for (SwitchInst::CaseIt i = SI->case_begin(), e = SI->case_end();
                i != e; ++i) {
             Constant *UnswitchValCandidate = i.getCaseValue();
             if (!BranchesInfo.isUnswitched(SI, UnswitchValCandidate)) {
               UnswitchVal = UnswitchValCandidate;
               break;
             }
           }
+        }
 
         if (!UnswitchVal)
           continue;
 
         if (UnswitchIfProfitable(LoopCond, UnswitchVal)) {
           ++NumSwitches;
+          // In case of a full LIV, UnswitchVal is the value we unswitched out.
+          // In case of a partial LIV, we only unswitch when its an AND-chain
+          // or OR-chain. In both cases switch input value simplifies to
+          // UnswitchVal.
+          BranchesInfo.setUnswitched(SI, UnswitchVal);
           return true;
         }
       }
     }
 
     // Scan the instructions to check for unswitchable values.
     for (BasicBlock::iterator BBI = (*I)->begin(), E = (*I)->end();
          BBI != E; ++BBI)
       if (SelectInst *SI = dyn_cast<SelectInst>(BBI)) {
         Value *LoopCond = FindLIVLoopCondition(SI->getCondition(),
-                                               currentLoop, Changed);
+                                               currentLoop, Changed).first;
         if (LoopCond && UnswitchIfProfitable(LoopCond,
                                              ConstantInt::getTrue(Context))) {
           ++NumSelects;
           return true;
         }
       }
   }
   return Changed;
bool LoopUnswitch::TryTrivialLoopUnswitch(bool &Changed) {
   BasicBlock *LoopExitBB = nullptr;
 
   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);
+                                           currentLoop, Changed).first;
 
     // 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
if (BranchInst *BI = dyn_cast<BranchInst>(CurrentTerm)) {
 
     UnswitchTrivialCondition(currentLoop, LoopCond, CondVal, LoopExitBB,
                              CurrentTerm);
     ++NumBranches;
     return true;
   } 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);
+                                           currentLoop, Changed).first;
 
     // 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;
 
     // Check to see if a successor of the switch is guaranteed to go to the
     // latch block or exit through a one exit block without having any
if (BranchInst *BI = dyn_cast<BranchInst>(CurrentTerm)) {
 
     // 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,
                              nullptr);
+
+    // We are only unswitching full LIV.
+    BranchesInfo.setUnswitched(SI, CondVal);
     ++NumSwitches;
     return true;
   }
   return false;
 }
 
 /// Split all of the edges from inside the loop to their exit blocks.
 /// Update the appropriate Phi nodes as we do so.
for (User *U : LIC->users()) {
 
     // TODO: We could do other simplifications, for example, turning
     // 'icmp eq LIC, Val' -> false.
 
     // If we know that LIC is not Val, use this info to simplify code.
     SwitchInst *SI = dyn_cast<SwitchInst>(UI);
     if (!SI || !isa<ConstantInt>(Val)) continue;
 
+    // NOTE: if a case value for the switch is unswitched out, we record it
+    // after the unswitch finishes. We can not record it here switch is not
+    // a direct user of the partial LIV.
     SwitchInst::CaseIt DeadCase = SI->findCaseValue(cast<ConstantInt>(Val));
     // Default case is live for multiple values.
     if (DeadCase == SI->case_default()) continue;
 
     // Found a dead case value.  Don't remove PHI nodes in the
     // successor if they become single-entry, those PHI nodes may
     // be in the Users list.
 
     BasicBlock *Switch = SI->getParent();
     BasicBlock *SISucc = DeadCase.getCaseSuccessor();
     BasicBlock *Latch = L->getLoopLatch();
 
-    BranchesInfo.setUnswitched(SI, Val);
-
     if (!SI->findCaseDest(SISucc)) continue;  // Edge is critical.
     // If the DeadCase successor dominates the loop latch, then the
     // transformation isn't safe since it will delete the sole predecessor edge
     // to the latch.
     if (Latch && DT->dominates(SISucc, Latch))
       continue;
 
     // FIXME: This is a hack.  We need to keep the successor around

test/Transforms/LoopUnswitch/basictest.ll

dec:
   br label %loop_begin
 default:
   br label %loop_exit
 loop_exit:
   ret i32 0
 ; CHECK: }
 }
 
+; Make sure we unswitch %a == 0 out of the loop.
+;
+; CHECK: define void @and_i2_as_switch_input(i2
+; CHECK: entry:
+; This is an indication that the loop has been unswitched.
+; CHECK: icmp eq i2 %a, 0
+; CHECK: br
+; There should be no more unswitching after the 1st unswitch.
+; CHECK-NOT: icmp eq
+; CHECK: ret
+define void @and_i2_as_switch_input(i2 %a) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i2 [ 0, %entry ], [ %inc, %for.inc ]
+  %and = and i2 %a, %i
+  %and1 = and i2 %and, %i
+  switch i2 %and1, label %sw.default [
+    i2 0, label %sw.bb
+    i2 1, label %sw.bb1
+  ]
+
+sw.bb:
+  br label %sw.epilog
+
+sw.bb1:
+  br label %sw.epilog
+
+sw.default:
+  br label %sw.epilog
+
+sw.epilog:
+  br label %for.inc
+
+for.inc:
+  %inc = add nsw i2 %i, 1
+  %cmp = icmp slt i2 %inc, 3 
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  ret void
+}
+
+; Make sure we unswitch %a == !0 out of the loop.
+;
+; CHECK: define void @or_i2_as_switch_input(i2
+; CHECK: entry:
+; This is an indication that the loop has been unswitched.
+; CHECK: icmp eq i2 %a, -1
+; CHECK: br
+; There should be no more unswitching after the 1st unswitch.
+; CHECK-NOT: icmp eq
+; CHECK: ret
+define void @or_i2_as_switch_input(i2 %a) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i2 [ 0, %entry ], [ %inc, %for.inc ]
+  %or = or i2 %a, %i
+  %or1 = or i2 %or, %i
+  switch i2 %or1, label %sw.default [
+    i2 2, label %sw.bb
+    i2 3, label %sw.bb1
+  ]
+
+sw.bb:
+  br label %sw.epilog
+
+sw.bb1:
+  br label %sw.epilog
+
+sw.default:
+  br label %sw.epilog
+
+sw.epilog:
+  br label %for.inc
+
+for.inc:
+  %inc = add nsw i2 %i, 1
+  %cmp = icmp slt i2 %inc, 3 
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  ret void
+}
+
+; Make sure we unswitch %a == !0 out of the loop. Even we do not
+; have it as a case value. Unswitching it out allows us to simplify
+; the or operator chain.
+;
+; CHECK: define void @or_i2_as_switch_input_unswitch_default(i2
+; CHECK: entry:
+; This is an indication that the loop has been unswitched.
+; CHECK: icmp eq i2 %a, -1
+; CHECK: br
+; There should be no more unswitching after the 1st unswitch.
+; CHECK-NOT: icmp eq
+; CHECK: ret
+define void @or_i2_as_switch_input_unswitch_default(i2 %a) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i2 [ 0, %entry ], [ %inc, %for.inc ]
+  %or = or i2 %a, %i
+  %or1 = or i2 %or, %i
+  switch i2 %or1, label %sw.default [
+    i2 1, label %sw.bb
+    i2 2, label %sw.bb1
+  ]
+
+sw.bb:
+  br label %sw.epilog
+
+sw.bb1:
+  br label %sw.epilog
+
+sw.default:
+  br label %sw.epilog
+
+sw.epilog:
+  br label %for.inc
+
+for.inc:
+  %inc = add nsw i2 %i, 1
+  %cmp = icmp slt i2 %inc, 3 
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  ret void
+}
+
+; Make sure we don't unswitch, as we can not find an input value %a
+; that will effectively unswitch 0 or 3 out of the loop.
+;
+; CHECK: define void @and_or_i2_as_switch_input(i2
+; CHECK: entry:
+; This is an indication that the loop has NOT been unswitched.
+; CHECK-NOT: icmp
+; CHECK: br
+define void @and_or_i2_as_switch_input(i2 %a) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i2 [ 0, %entry ], [ %inc, %for.inc ]
+  %and = and i2 %a, %i 
+  %or = or i2 %and, %i
+  switch i2 %or, label %sw.default [
+    i2 0, label %sw.bb
+    i2 3, label %sw.bb1
+  ]
+
+sw.bb:
+  br label %sw.epilog
+
+sw.bb1:
+  br label %sw.epilog
+
+sw.default:
+  br label %sw.epilog
+
+sw.epilog:
+  br label %for.inc
+
+for.inc:
+  %inc = add nsw i2 %i, 1
+  %cmp = icmp slt i2 %inc, 3 
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  ret void
+}
+
+; Make sure we don't unswitch, as we can not find an input value %a
+; that will effectively unswitch true/false out of the loop.
+;
+; CHECK: define void @and_or_i1_as_branch_input(i1
+; CHECK: entry:
+; This is an indication that the loop has NOT been unswitched.
+; CHECK-NOT: icmp
+; CHECK: br
+define void @and_or_i1_as_branch_input(i1 %a) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i1 [ 0, %entry ], [ %inc, %for.inc ]
+  %and = and i1 %a, %i 
+  %or = or i1 %and, %i
+  br i1 %or, label %sw.bb, label %sw.bb1
+
+sw.bb:
+  br label %sw.epilog
+
+sw.bb1:
+  br label %sw.epilog
+
+sw.epilog:
+  br label %for.inc
+
+for.inc:
+  %inc = add nsw i1 %i, 1
+  %cmp = icmp slt i1 %inc, 1 
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  ret void
+}
 
 declare void @incf() noreturn
 declare void @decf() noreturn
 declare void @conv() convergent