[PM/LoopUnswitch] Support partial trivial unswitching.

[chandlerc]

Summary

The idea of partial unswitching is to take a *part* of a branch's
condition that is loop invariant and just unswitching that part. This
primarily makes sense with i1 conditions of branches as opposed to
switches. When dealing with i1 conditions, we can easily extract loop
invariant inputs to a a branch and unswitch them to test them entirely
outside the loop.

As part of this, we now create much more significant cruft in the loop
body, so this relies on adding cleanup passes to the loop pipeline and
revisiting unswitched loops to do that cleanup before continuing to
process them.

This already appears to be more powerful at unswitching than the old
loop unswitch pass, and so I'd appreciate pretty careful review in case
I'm just missing some correctness checks. The LIV-loop-condition test
case is not unswitched by the old unswitch pass, but is with this pass.

Depends on D47408.

[fedor.sergeev]

What about finally renaming SimpleLoopUnswitch into something "less simple"? :)

[chandlerc]

OMG yes I want to rename it. I mostly have been avoiding this because i
have too many patches stacked on top of one another.

That said, if you have a better name, I'm happy to prioritize fixnig the
name. No arguments about it needing to happen. I also just don't really
know what to call it....

[sanjoy]

llvm/include/llvm/Analysis/Utils/Local.h

144 ↗

(On Diff #146156)

Why not s/provided vector/DeadInsts?

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

76 ↗

(On Diff #146156)

Nit: wrap?

322 ↗

(On Diff #146156)

Stray line?

333 ↗

(On Diff #146156)

Would be nice to be a bit more specific "For partial trivial unswitching of a condition".

398 ↗

(On Diff #146156)

Might be nice to assert here that if LoopExitSuccIdx == 0 then cast<Instruction>(BI.getCondition())->getOpcode() is Or etc.

llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll

511 ↗

(On Diff #146156)

I think we need to test a few more cases here:

• Branch on chain of ands
• Branch on mix of ands and ors and possibly other operations

[sanjoy]

Actually now that I think of it, I wonder if looking back through the condition expression tree to find all loop invariant values is necessary -- if we see these kinds of cases (loop_varying & loop_invariant0) & loop_invariant1) then perhaps we should teach LICM et. al. to reassociate and hoist the loop_invariant0 & loop_invariant1 bit instead of worrying about them in loop unswitch?

[llvm-commits]

• msg-29308-243.txt162 BDownload

[chandlerc]

Update with fixes from code review.

[chandlerc]

Thanks for the review!

llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll

511 ↗

(On Diff #146156)

So, we test branching on an and above. There is no different logic in handling N and instructions vs. N or instructions, so I didn't add that test as it didn't seem to add value compared to these two.

I can definitely mix some non or operations into this chain. Would that be enough coverage? Just trying to understand what the goal is of the added tests.

While I was adding these tests, it actually exposed a weakness in the instruction simplification we do here. Fixing it proved... a bit tricky. I've added somewhat complex logic so that when we need to simplify multiple different invariants we can correctly simplify defs before uses. However, now I need to add a better test to cover the nasty case for simplifying -- when we need to simplify around a diamond in the CFG. I'm going ahead and uploading the patch while I craft a test case for that so you can yell if I'm making this way harder than it needs to be.

[sanjoy]

llvm/include/llvm/Analysis/Utils/Local.h

144 ↗

(On Diff #146156)

Wasn't done?

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

208 ↗

(On Diff #148363)

"just walk"

228 ↗

(On Diff #148363)

Maybe this can go in a different "loop instruction simplify" pass? Just above you say "the only real goal of this is to do very basic cleanup of unswitched conditions. We don't need powerful tools here. A proper pass can be scheduled to do more comprehensive cleanup." :)

I'm worried that doing all of this in a single pass makes it less granular and harder to test -- we only get to the see the IR after all of this cross-block simplification has happened, which means it is harder to figure out what exactly loop unswitching did.

llvm/lib/Transforms/Utils/Local.cpp

447 ↗

(On Diff #148363)

Can we end up doing a double free here? Like

%a = add %p, 1
%b = add %a, 1

and DeadInsts is [%a,%b] -- we'll add %a back to the worklist after visiting %b and then eraseFromParent %a twice.

We could add a precondition here to avoid this, but it seems better to just handle this case.

453 ↗

(On Diff #148363)

I think OpU.set(nullptr) is more readable here. Otherwise this reads like we're just modifying a local variable.

llvm/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll

511 ↗

(On Diff #146156)

Just trying to understand what the goal is of the added tests.

Mostly just making sure that the logic to "gather" the various loop invariant (transitive) operands keeps working as intended.

[fedor.sergeev]

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

228 ↗

(On Diff #148363)

Second that.
With our current LoopUnswitch use we tend to run massive cleanup after it, and I'm sure we will continue doing that with the new LoopUnswitch. Having a small cleanup in LoopUnswitch itself will not do a full job, and as Sanjoy rightfully points out it might make IR investigations a bit more complicated.

[chandlerc]

Ok, after some discussion, I'm pretty convinced we just need to stop simplifying here.

I was hoping to do at least *some* simplification, but there really is no good middle ground and so I think its better to *completely* nuke this.

Unfortunately that is going to require ... some more work.

1. I think the current pass pipeline doesn't have a good cleanup pass scheduled. I don't think we even have the *right* cleanup pass here, but I'll poke around at some of the nascent ones. So first step will be that I need to go build / finish / polish a good instruction cleanup loop pass. And then I'll need to add it to the main pipelines. Prepare for patches to that effect.

2. I'll then need to do the removal of all simplifying here in a separate patch and update testing to reflect it. Specifically, when we *only* do trivial unswitching, we were falling through to try non-trivial unswitching immediately. We will have to not do that, and instead re-add the now mutated loop to the pass manager to re-visit (much like we do with non-trivial unswitching) and rely on it then iterating for us. This will ... not be a somewhat surprisingly significant behavior change. I think its good, but its worth noting. This will essentially make loop-unswitch a fixed-point pass in the pass pipeline, but it will do so using the pass manager. I'm not aware of any other passes that currently do this. Anyways, brave new world. The pipeline will now be ... truly dynamic.

3. Then I can return to this patch and land it w/o added simplification.

Make sense?

I've updated the code here to reflect what I think it will end up looking like. I can't update the tests yet, will do that when I get back to #3.

I'll also probably pull the recursive deletion API into one of the patches in #1 -- would you be OK with that going in on its own? I don't have a good way of testing it that way though. Otherwise, I'll let it float around and try to figure out when/if I need it.

-Chandler

llvm/include/llvm/Analysis/Utils/Local.h

144 ↗

(On Diff #146156)

Oh, sorry. I changed it in the last paragraph, but not the top one.

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

228 ↗

(On Diff #148363)

After a bunch of discussion, there really isn't even a good middle ground here.

I think we have to rip *all* of the cleanup out. This is going to make things harder to do as it changes how the current pass behaves. See my top-level response.

llvm/lib/Transforms/Utils/Local.cpp

447 ↗

(On Diff #148363)

This is already in the precondition documentation... Specifically that the instructions in the vector must have no uses.

I've added an assert to ensure this holds. I'm happy to widen the contract if you think users would benefit from it though.

[chandlerc]

Updated patch just to show where the code ended up.

[hfinkel]

and instead re-add the now mutated loop to the pass manager to re-visit (much like we do with non-trivial unswitching) and rely on it then iterating for us. This will ... not be a somewhat surprisingly significant behavior change. I think its good, but its worth noting. This will essentially make loop-unswitch a fixed-point pass in the pass pipeline, but it will do so using the pass manager. I'm not aware of any other passes that currently do this. Anyways, brave new world. The pipeline will now be ... truly dynamic.

That's really nice :-)

[fedor.sergeev]

The pipeline will now be ... truly dynamic.

Hoorray! :)

Make sense?

sounds good to me

[chandlerc]

Update to rebase on the patch that adds loop-instsimplify and loop-simplifycfg
to the loop pass pipeline as cleanup passes to minimize the cleanup necessary
in the unswitch pass itself.

Also includes a bugfix found when testing this on the llvm test-suite, and new
test cases to cover this area of code.

[fedor.sergeev]

Not exactly related to this set of changes...
I just have discovered that old LoopUnswitch performs SCEV cache invalidation when it does nontrivial updates to CFG
(forgetLoop in LoopUnswitch::unswitchNontrivialCondition) .

Perhaps makes sense to do the same in new loop unswitch as well?

[chandlerc]

In D46706#1113840, @fedor.sergeev wrote:

Not exactly related to this set of changes...
I just have discovered that old LoopUnswitch performs SCEV cache invalidation when it does nontrivial updates to CFG
(forgetLoop in LoopUnswitch::unswitchNontrivialCondition) .

Perhaps makes sense to do the same in new loop unswitch as well?

Yes, will attack that in a separate patch unless someone gets there first.

[chandlerc]

Rebase now that all the loop-instsimplify stuff has landed.

This should be ready for another round of review?

[chandlerc]

Rebase and ping.

[chandlerc]

Rebase and ping yet again. Would really like to get a review here...

[chandlerc]

Crickets.....

[sanjoy]

This lgtm.

Have you considered making full unswitching a special case for partial unswitching? It seems like we could "partially unswitch" br %loop_invariant, label %t, label %f into br true/false, label %t, label %f and have a beefed up LoopSimplifyCFG clean this up?

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

84 ↗

(On Diff #150040)

Assert that Root is loop invariant? That justifies why we're not adding it to Invariants.

[chandlerc]

In D46706#1127667, @sanjoy wrote:

This lgtm.

Have you considered making full unswitching a special case for partial unswitching? It seems like we could "partially unswitch" br %loop_invariant, label %t, label %f into br true/false, label %t, label %f and have a beefed up LoopSimplifyCFG clean this up?

We'd have to make LoopSimplifyCFG substantially more powerful, and all of the complex CFG-fixing logic here would need to be ported to it.

We can do that, but I'm not sure how valuable it is in practice. The really tricky thing is that we'd have to have LoopSimplifyCFG trigger the iteration when it re-shapes the loop nest. That seems a bit more clear here at the moment... but happy to chat about it further and if it makes sense, we can move all the CFG rewriting logic to LoopSimplifyCFG and make that change.

[reames]

Sorry for being late to the party, but a couple of optional post commit style comments for possible follow up. Nothing major, just ideas on how to share code and reduce a possible ordering sensitivity.

llvm/trunk/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

81

Hm, with some slight generalization, this could be a generally useful utility elsewhere. This is basically handling reduction operations spelled with multiple instructions. Seems like that's a common enough pattern (say, instcombine?) to be worth drawing out into a generic visitReductionOperands(Root, FilterFunc)?

248

I believe you've got an pass ordering sensitivity here which isn't present in the old pass. The old pass deliberately uses "makeLoopInvariant" which is like isLoopInvariant, but it will hoist trivially hoistable instructions as well. By not doing this, the new pass will be very sensitive to having ideal (i.e. fully LICMed) IR. This may be problematic if any other loop pass is run between LICM and the new unswitch.

349

Hm, possibly out of scope, but I'll mention it anyway.

This ends up producing a reduce chain rather than a tree. We're going to end up trying to reassociate that later. Is it worthwhile having an interface which produces an appropriate tree to start with? Maybe this is an API that make sense on IR Builder? Something like CreateOr(ArrayRef<Value> Ops)?

377

Hm, we have this pattern all over. Might be time for a getBoolean(bool) on ConstantInt?

[chandlerc]

In D46706#1138515, @reames wrote:

Sorry for being late to the party, but a couple of optional post commit style comments for possible follow up. Nothing major, just ideas on how to share code and reduce a possible ordering sensitivity.

No problem. Mostly following up here on the more discussion-y points.

llvm/trunk/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

248

Yes, this has been true throughout the new pass. IMO, it makes the pass quite a bit simpler as we get the invariant that if we didn't unswitch something, we didn't mutate the IR. And that matters a lot more in the new PM with caching of analyses.

I'd rather fix the pass pipelines to be much more careful about running LICM in the appropriate places. So far we haven't hit any issues, but it is definitely something that may come up in the future.

Does that seem reasonable?

349

I think we canonicalize to chains pretty commonly to make analysis easier anyways, and rely on later passes forming trees where beneficial (most places).

llvm/trunk/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/CodeMetrics.h"
+#include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopIterator.h"
 #include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/Utils/Local.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
static cl::opt<bool> EnableNonTrivialUnswitch(
     "enable-nontrivial-unswitch", cl::init(false), cl::Hidden,
     cl::desc("Forcibly enables non-trivial loop unswitching rather than "
              "following the configuration passed into the pass."));
 
 static cl::opt<int>
     UnswitchThreshold("unswitch-threshold", cl::init(50), cl::Hidden,
                       cl::desc("The cost threshold for unswitching a loop."));
 
-static void replaceLoopUsesWithConstant(Loop &L, Value &LIC,
-                                        Constant &Replacement) {
-  assert(!isa<Constant>(LIC) && "Why are we unswitching on a constant?");
+/// Collect all of the loop invariant input values transitively used by the
+/// homogeneous instruction graph from a given root.
+///
+/// This essentially walks from a root recursively through loop variant operands
+/// which have the exact same opcode and finds all inputs which are loop
+/// invariant. For some operations these can be re-associated and unswitched out
+/// of the loop entirely.
+static SmallVector<Value *, 4>
+collectHomogenousInstGraphLoopInvariants(Loop &L, Instruction &Root,
+                                         LoopInfo &LI) {
+  SmallVector<Value *, 4> Invariants;
+  assert(!L.isLoopInvariant(&Root) &&
+         "Only need to walk the graph if root itself is not invariant.");
+
+  // Build a worklist and recurse through operators collecting invariants.
+  SmallVector<Instruction *, 4> Worklist;
+  SmallPtrSet<Instruction *, 8> Visited;
+  Worklist.push_back(&Root);
+  Visited.insert(&Root);
+  do {
+    Instruction &I = *Worklist.pop_back_val();
+    for (Value *OpV : I.operand_values()) {
+      // Skip constants as unswitching isn't interesting for them.
+      if (isa<Constant>(OpV))
+        continue;
+
+      // Add it to our result if loop invariant.
+      if (L.isLoopInvariant(OpV)) {
+        Invariants.push_back(OpV);
+        continue;
+      }
+
+      // If not an instruction with the same opcode, nothing we can do.
+      Instruction *OpI = dyn_cast<Instruction>(OpV);
+      if (!OpI || OpI->getOpcode() != Root.getOpcode())
+        continue;
+
+      // Visit this operand.
+      if (Visited.insert(OpI).second)
+        Worklist.push_back(OpI);
+    }
+  } while (!Worklist.empty());
+
+  return Invariants;
+}
+
+static void replaceLoopInvariantUses(Loop &L, Value *Invariant,
+  assert(!isa<Constant>(Invariant) && "Why are we unswitching on a constant?");
 
   // Replace uses of LIC in the loop with the given constant.
-  for (auto UI = LIC.use_begin(), UE = LIC.use_end(); UI != UE;) {
+  for (auto UI = Invariant->use_begin(), UE = Invariant->use_end(); UI != UE;) {
     // Grab the use and walk past it so we can clobber it in the use list.
     Use *U = &*UI++;
     Instruction *UserI = dyn_cast<Instruction>(U->getUser());
-    if (!UserI || !L.contains(UserI))
-      continue;
 
     // Replace this use within the loop body.
-    *U = &Replacement;
+    if (UserI && L.contains(UserI))
+      U->set(&Replacement);
   }
 }
 
 /// Check that all the LCSSA PHI nodes in the loop exit block have trivial
 /// incoming values along this edge.
 static bool areLoopExitPHIsLoopInvariant(Loop &L, BasicBlock &ExitingBB,
                                          BasicBlock &ExitBB) {
   for (Instruction &I : ExitBB) {

 ///
 /// Because the exit block remains an exit from the loop, this rewrites the
 /// LCSSA PHI nodes in it to remove the unswitched edge and introduces PHI
 /// nodes into the unswitched basic block to select between the value in the
 /// old preheader and the loop exit.
 static void rewritePHINodesForExitAndUnswitchedBlocks(BasicBlock &ExitBB,
                                                       BasicBlock &UnswitchedBB,
                                                       BasicBlock &OldExitingBB,
-                                                      BasicBlock &OldPH) {
+                                                      BasicBlock &OldPH,
+                                                      bool FullUnswitch) {
   assert(&ExitBB != &UnswitchedBB &&
          "Must have different loop exit and unswitched blocks!");
   Instruction *InsertPt = &*UnswitchedBB.begin();
   for (PHINode &PN : ExitBB.phis()) {
     auto *NewPN = PHINode::Create(PN.getType(), /*NumReservedValues*/ 2,
                                   PN.getName() + ".split", InsertPt);
 
     // Walk backwards over the old PHI node's inputs to minimize the cost of
     // removing each one. We have to do this weird loop manually so that we
     // create the same number of new incoming edges in the new PHI as we expect
     // each case-based edge to be included in the unswitched switch in some
     // cases.
     // FIXME: This is really, really gross. It would be much cleaner if LLVM
     // allowed us to create a single entry for a predecessor block without
     // having separate entries for each "edge" even though these edges are
     // required to produce identical results.
     for (int i = PN.getNumIncomingValues() - 1; i >= 0; --i) {
       if (PN.getIncomingBlock(i) != &OldExitingBB)
         continue;
 
-      Value *Incoming = PN.removeIncomingValue(i);
+      Value *Incoming = PN.getIncomingValue(i);
+      if (FullUnswitch)
+        // No more edge from the old exiting block to the exit block.
+        PN.removeIncomingValue(i);
+
       NewPN->addIncoming(Incoming, &OldPH);
     }
 
     // Now replace the old PHI with the new one and wire the old one in as an
     // input to the new one.
     PN.replaceAllUsesWith(NewPN);
     NewPN->addIncoming(&PN, &ExitBB);
   }

 /// (splitting the exit block as necessary). It simplifies the branch within
 /// the loop to an unconditional branch but doesn't remove it entirely. Further
 /// cleanup can be done with some simplify-cfg like pass.
 static bool unswitchTrivialBranch(Loop &L, BranchInst &BI, DominatorTree &DT,
                                   LoopInfo &LI) {
   assert(BI.isConditional() && "Can only unswitch a conditional branch!");
   LLVM_DEBUG(dbgs() << "  Trying to unswitch branch: " << BI << "\n");
 
-  Value *LoopCond = BI.getCondition();
+  // The loop invariant values that we want to unswitch.
+  SmallVector<Value *, 4> Invariants;
 
-  // Need a trivial loop condition to unswitch.
-  if (!L.isLoopInvariant(LoopCond))
+  // When true, we're fully unswitching the branch rather than just unswitching
+  // some input conditions to the branch.
+  bool FullUnswitch = false;
+
+  if (L.isLoopInvariant(BI.getCondition())) {
+    Invariants.push_back(BI.getCondition());
+    FullUnswitch = true;
+  } else {
+    if (auto *CondInst = dyn_cast<Instruction>(BI.getCondition()))
+      Invariants = collectHomogenousInstGraphLoopInvariants(L, *CondInst, LI);
+    if (Invariants.empty())
+      // Couldn't find invariant inputs!
       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.
-  ConstantInt *CondVal = ConstantInt::getTrue(BI.getContext());
-  ConstantInt *Replacement = ConstantInt::getFalse(BI.getContext());
+  // Check that one of the branch's successors exits, and which one.
+  bool ExitDirection = true;
   int LoopExitSuccIdx = 0;
   auto *LoopExitBB = BI.getSuccessor(0);
   if (L.contains(LoopExitBB)) {
-    std::swap(CondVal, Replacement);
+    ExitDirection = false;
     LoopExitSuccIdx = 1;
     LoopExitBB = BI.getSuccessor(1);
     if (L.contains(LoopExitBB))
       return false;
   }
   auto *ContinueBB = BI.getSuccessor(1 - LoopExitSuccIdx);
   auto *ParentBB = BI.getParent();
   if (!areLoopExitPHIsLoopInvariant(L, *ParentBB, *LoopExitBB))
     return false;
 
-  LLVM_DEBUG(dbgs() << "    unswitching trivial branch when: " << CondVal
-                    << " == " << LoopCond << "\n");
+  // When unswitching only part of the branch's condition, we need the exit
+  // block to be reached directly from the partially unswitched input. This can
+  // be done when the exit block is along the true edge and the branch condition
+  // is a graph of `or` operations, or the exit block is along the false edge
+  // and the condition is a graph of `and` operations.
+  if (!FullUnswitch) {
+    if (ExitDirection) {
+      if (cast<Instruction>(BI.getCondition())->getOpcode() != Instruction::Or)
+        return false;
+    } else {
+      if (cast<Instruction>(BI.getCondition())->getOpcode() != Instruction::And)
+        return false;
+    }
+  }
+
+  LLVM_DEBUG({
+    dbgs() << "    unswitching trivial invariant conditions for: " << BI
+           << "\n";
+    for (Value *Invariant : Invariants) {
+      dbgs() << "      " << *Invariant << " == true";
+      if (Invariant != Invariants.back())
+        dbgs() << " ||";
+      dbgs() << "\n";
+    }
+  });
 
   // Split the preheader, so that we know that there is a safe place to insert
   // the conditional branch. We will change the preheader to have a conditional
   // branch on LoopCond.
   BasicBlock *OldPH = L.getLoopPreheader();
   BasicBlock *NewPH = SplitEdge(OldPH, L.getHeader(), &DT, &LI);
 
   // Now that we have a place to insert the conditional branch, create a place
   // to branch to: this is the exit block out of the loop that we are
   // unswitching. We need to split this if there are other loop predecessors.
   // Because the loop is in simplified form, *any* other predecessor is enough.
   BasicBlock *UnswitchedBB;
-  if (BasicBlock *PredBB = LoopExitBB->getUniquePredecessor()) {
-    (void)PredBB;
-    assert(PredBB == BI.getParent() &&
+  if (FullUnswitch && LoopExitBB->getUniquePredecessor()) {
+    assert(LoopExitBB->getUniquePredecessor() == BI.getParent() &&
            "A branch's parent isn't a predecessor!");
     UnswitchedBB = LoopExitBB;
   } else {
     UnswitchedBB = SplitBlock(LoopExitBB, &LoopExitBB->front(), &DT, &LI);
   }
 
-  // Now splice the branch to gate reaching the new preheader and re-point its
-  // successors.
-  OldPH->getInstList().splice(std::prev(OldPH->end()),
-                              BI.getParent()->getInstList(), BI);
+  // Actually move the invariant uses into the unswitched position. If possible,
+  // we do this by moving the instructions, but when doing partial unswitching
+  // we do it by building a new merge of the values in the unswitched position.
   OldPH->getTerminator()->eraseFromParent();
+  if (FullUnswitch) {
+    // If fully unswitching, we can use the existing branch instruction.
+    // Splice it into the old PH to gate reaching the new preheader and re-point
+    // its successors.
+    OldPH->getInstList().splice(OldPH->end(), BI.getParent()->getInstList(),
+                                BI);
     BI.setSuccessor(LoopExitSuccIdx, UnswitchedBB);
     BI.setSuccessor(1 - LoopExitSuccIdx, NewPH);
 
     // Create a new unconditional branch that will continue the loop as a new
     // terminator.
     BranchInst::Create(ContinueBB, ParentBB);
+  } else {
+    // Only unswitching a subset of inputs to the condition, so we will need to
+    // build a new branch that merges the invariant inputs.
+    IRBuilder<> IRB(OldPH);
+    Value *Cond = Invariants.front();
+    if (ExitDirection)
+      assert(cast<Instruction>(BI.getCondition())->getOpcode() ==
+                 Instruction::Or &&
+             "Must have an `or` of `i1`s for the condition!");
+    else
+      assert(cast<Instruction>(BI.getCondition())->getOpcode() ==
+                 Instruction::And &&
+             "Must have an `and` of `i1`s for the condition!");
+    for (Value *Invariant :
+         make_range(std::next(Invariants.begin()), Invariants.end()))
+      if (ExitDirection)
+        Cond = IRB.CreateOr(Cond, Invariant);
+      else
+        Cond = IRB.CreateAnd(Cond, Invariant);
+
+    BasicBlock *Succs[2];
+    Succs[LoopExitSuccIdx] = UnswitchedBB;
+    Succs[1 - LoopExitSuccIdx] = NewPH;
+    IRB.CreateCondBr(Cond, Succs[0], Succs[1]);
+  }
 
   // Rewrite the relevant PHI nodes.
   if (UnswitchedBB == LoopExitBB)
     rewritePHINodesForUnswitchedExitBlock(*UnswitchedBB, *ParentBB, *OldPH);
   else
     rewritePHINodesForExitAndUnswitchedBlocks(*LoopExitBB, *UnswitchedBB,
-                                              *ParentBB, *OldPH);
+                                              *ParentBB, *OldPH, FullUnswitch);
 
   // Now we need to update the dominator tree.
-  DT.applyUpdates(
-      {{DT.Delete, ParentBB, UnswitchedBB}, {DT.Insert, OldPH, UnswitchedBB}});
+  DT.insertEdge(OldPH, UnswitchedBB);
+  if (FullUnswitch)
+    DT.deleteEdge(ParentBB, UnswitchedBB);
+
+  // The constant we can replace all of our invariants with inside the loop
+  // body. If any of the invariants have a value other than this the loop won't
+  // be entered.
+  ConstantInt *Replacement = ExitDirection
+                                 ? ConstantInt::getFalse(BI.getContext())
+                                 : ConstantInt::getTrue(BI.getContext());
 
   // Since this is an i1 condition we can also trivially replace uses of it
   // within the loop with a constant.
-  replaceLoopUsesWithConstant(L, *LoopCond, *Replacement);
+  for (Value *Invariant : Invariants)
+    replaceLoopInvariantUses(L, Invariant, *Replacement);
 
   ++NumTrivial;
   ++NumBranches;
   return true;
 }
 
 /// Unswitch a trivial switch if the condition is loop invariant.
 ///
static bool unswitchTrivialSwitch(Loop &L, SwitchInst &SI, DominatorTree &DT,
   // ranges aren't quite powerful enough yet.
   if (DefaultExitBB) {
     if (pred_empty(DefaultExitBB)) {
       UnswitchedExitBBs.insert(DefaultExitBB);
       rewritePHINodesForUnswitchedExitBlock(*DefaultExitBB, *ParentBB, *OldPH);
     } else {
       auto *SplitBB =
           SplitBlock(DefaultExitBB, &DefaultExitBB->front(), &DT, &LI);
-      rewritePHINodesForExitAndUnswitchedBlocks(*DefaultExitBB, *SplitBB,
-                                                *ParentBB, *OldPH);
+      rewritePHINodesForExitAndUnswitchedBlocks(
+          *DefaultExitBB, *SplitBB, *ParentBB, *OldPH, /*FullUnswitch*/ true);
       DefaultExitBB = SplitExitBBMap[DefaultExitBB] = SplitBB;
     }
   }
   // Note that we must use a reference in the for loop so that we update the
   // container.
   for (auto &CasePair : reverse(ExitCases)) {
     // Grab a reference to the exit block in the pair so that we can update it.
     BasicBlock *ExitBB = CasePair.second;
 
     // If this case is the last edge into the exit block, we can simply reuse it
     // as it will no longer be a loop exit. No mapping necessary.
     if (pred_empty(ExitBB)) {
       // Only rewrite once.
       if (UnswitchedExitBBs.insert(ExitBB).second)
         rewritePHINodesForUnswitchedExitBlock(*ExitBB, *ParentBB, *OldPH);
       continue;
     }
 
     // Otherwise we need to split the exit block so that we retain an exit
     // block from the loop and a target for the unswitched condition.
     BasicBlock *&SplitExitBB = SplitExitBBMap[ExitBB];
     if (!SplitExitBB) {
       // If this is the first time we see this, do the split and remember it.
       SplitExitBB = SplitBlock(ExitBB, &ExitBB->front(), &DT, &LI);
-      rewritePHINodesForExitAndUnswitchedBlocks(*ExitBB, *SplitExitBB,
-                                                *ParentBB, *OldPH);
+      rewritePHINodesForExitAndUnswitchedBlocks(
+          *ExitBB, *SplitExitBB, *ParentBB, *OldPH, /*FullUnswitch*/ true);
     }
     // Update the case pair to point to the split block.
     CasePair.second = SplitExitBB;
   }
 
   // Now add the unswitched cases. We do this in reverse order as we built them
   // in reverse order.
   for (auto CasePair : reverse(ExitCases)) {
do {
     // Found a trivial condition candidate: non-foldable conditional branch. If
     // we fail to unswitch this, we can't do anything else that is trivial.
     if (!unswitchTrivialBranch(L, *BI, DT, LI))
       return Changed;
 
     // Mark that we managed to unswitch something.
     Changed = true;
 
-    // We unswitched the branch. This should always leave us with an
-    // unconditional branch that we can follow now.
+    // If we only unswitched some of the conditions feeding the branch, we won't
+    // have collapsed it to a single successor.
     BI = cast<BranchInst>(CurrentBB->getTerminator());
-    assert(!BI->isConditional() &&
-           "Cannot form a conditional branch by unswitching1");
+    if (BI->isConditional())
+      return Changed;
+
+    // Follow the newly unconditional branch into its successor.
     CurrentBB = BI->getSuccessor(0);
 
     // When continuing, if we exit the 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.
   } while (L.contains(CurrentBB) && Visited.insert(CurrentBB).second);
 
   return Changed;
for (auto *ClonedBB : ClonedLoopBlocks)
     if (!BlocksInClonedLoop.count(ClonedBB))
       UnloopedBlockSet.insert(ClonedBB);
 
   // Copy the cloned exits and sort them in ascending loop depth, we'll work
   // backwards across these to process them inside out. The order shouldn't
   // matter as we're just trying to build up the map from inside-out; we use
   // the map in a more stably ordered way below.
   auto OrderedClonedExitsInLoops = ClonedExitsInLoops;
-  llvm::sort(OrderedClonedExitsInLoops.begin(),
-             OrderedClonedExitsInLoops.end(),
+  llvm::sort(OrderedClonedExitsInLoops.begin(), OrderedClonedExitsInLoops.end(),
              [&](BasicBlock *LHS, BasicBlock *RHS) {
                return ExitLoopMap.lookup(LHS)->getLoopDepth() <
                       ExitLoopMap.lookup(RHS)->getLoopDepth();
              });
 
   // Populate the existing ExitLoopMap with everything reachable from each
   // exit, starting from the inner most exit.
   while (!UnloopedBlockSet.empty() && !OrderedClonedExitsInLoops.empty()) {

llvm/trunk/test/Transforms/SimpleLoopUnswitch/LIV-loop-condtion.ll

 ; RUN: opt < %s -simple-loop-unswitch -S 2>&1 | FileCheck %s
 
 ; This is to test trivial loop unswitch only happens when trivial condition
 ; itself is an LIV loop condition (not partial LIV which could occur in and/or).
 
 define i32 @test(i1 %cond1, i32 %var1) {
+; CHECK-LABEL: define i32 @test(
 entry:
   br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 %cond1, label %entry.split, label %loop_exit.split
+;
+; CHECK:       entry.split:
+; CHECK-NEXT:    br label %loop_begin
 
 loop_begin:
   %var3 = phi i32 [%var1, %entry], [%var2, %do_something]
   %cond2 = icmp eq i32 %var3, 10
   %cond.and = and i1 %cond1, %cond2
-  
-; %cond.and only has %cond1 as LIV so no unswitch should happen.
-; CHECK: br i1 %cond.and, label %do_something, label %loop_exit
-  br i1 %cond.and, label %do_something, label %loop_exit 
+  br i1 %cond.and, label %do_something, label %loop_exit
+; CHECK:       loop_begin:
+; CHECK-NEXT:    %[[VAR3:.*]] = phi i32
+; CHECK-NEXT:    %[[COND2:.*]] = icmp eq i32 %[[VAR3]], 10
+; CHECK-NEXT:    %[[COND_AND:.*]] = and i1 true, %[[COND2]]
+; CHECK-NEXT:    br i1 %[[COND_AND]], label %do_something, label %loop_exit
 
 do_something:
   %var2 = add i32 %var3, 1
   call void @some_func() noreturn nounwind
   br label %loop_begin
 
 loop_exit:
   ret i32 0
 }
 
 declare void @some_func() noreturn 

llvm/trunk/test/Transforms/SimpleLoopUnswitch/trivial-unswitch.ll

 ; CHECK:       if.then23:
 ; CHECK-NEXT:    br label %cleanup
 
 cleanup:
   ret void
 ; CHECK:       cleanup:
 ; CHECK-NEXT:    ret void
 }
+
+define i32 @test_partial_condition_unswitch_and(i32* %var, i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test_partial_condition_unswitch_and(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 %cond1, label %entry.split, label %loop_exit.split
+;
+; CHECK:       entry.split:
+; CHECK-NEXT:    br i1 %cond2, label %entry.split.split, label %loop_exit
+;
+; CHECK:       entry.split.split:
+; CHECK-NEXT:    br label %loop_begin
+
+loop_begin:
+  br i1 %cond1, label %continue, label %loop_exit
+; CHECK:       loop_begin:
+; CHECK-NEXT:    br label %continue
+
+continue:
+  %var_val = load i32, i32* %var
+  %var_cond = trunc i32 %var_val to i1
+  %cond_and = and i1 %var_cond, %cond2
+  br i1 %cond_and, label %do_something, label %loop_exit
+; CHECK:       continue:
+; CHECK-NEXT:    %[[VAR:.*]] = load i32
+; CHECK-NEXT:    %[[VAR_COND:.*]] = trunc i32 %[[VAR]] to i1
+; CHECK-NEXT:    %[[COND_AND:.*]] = and i1 %[[VAR_COND]], true
+; CHECK-NEXT:    br i1 %[[COND_AND]], label %do_something, label %loop_exit
+
+do_something:
+  call void @some_func() noreturn nounwind
+  br label %loop_begin
+; CHECK:       do_something:
+; CHECK-NEXT:    call
+; CHECK-NEXT:    br label %loop_begin
+
+loop_exit:
+  ret i32 0
+; CHECK:       loop_exit:
+; CHECK-NEXT:    br label %loop_exit.split
+;
+; CHECK:       loop_exit.split:
+; CHECK-NEXT:    ret
+}
+
+define i32 @test_partial_condition_unswitch_or(i32* %var, i1 %cond1, i1 %cond2, i1 %cond3, i1 %cond4, i1 %cond5, i1 %cond6) {
+; CHECK-LABEL: @test_partial_condition_unswitch_or(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %[[INV_OR1:.*]] = or i1 %cond4, %cond2
+; CHECK-NEXT:    %[[INV_OR2:.*]] = or i1 %[[INV_OR1]], %cond3
+; CHECK-NEXT:    %[[INV_OR3:.*]] = or i1 %[[INV_OR2]], %cond1
+; CHECK-NEXT:    br i1 %[[INV_OR3]], label %loop_exit.split, label %entry.split
+;
+; CHECK:       entry.split:
+; CHECK-NEXT:    br label %loop_begin
+
+loop_begin:
+  %var_val = load i32, i32* %var
+  %var_cond = trunc i32 %var_val to i1
+  %cond_or1 = or i1 %var_cond, %cond1
+  %cond_or2 = or i1 %cond2, %cond3
+  %cond_or3 = or i1 %cond_or1, %cond_or2
+  %cond_xor1 = xor i1 %cond5, %var_cond
+  %cond_and1 = and i1 %cond6, %var_cond
+  %cond_or4 = or i1 %cond_xor1, %cond_and1
+  %cond_or5 = or i1 %cond_or3, %cond_or4
+  %cond_or6 = or i1 %cond_or5, %cond4
+  br i1 %cond_or6, label %loop_exit, label %do_something
+; CHECK:       loop_begin:
+; CHECK-NEXT:    %[[VAR:.*]] = load i32
+; CHECK-NEXT:    %[[VAR_COND:.*]] = trunc i32 %[[VAR]] to i1
+; CHECK-NEXT:    %[[COND_OR1:.*]] = or i1 %[[VAR_COND]], false
+; CHECK-NEXT:    %[[COND_OR2:.*]] = or i1 false, false
+; CHECK-NEXT:    %[[COND_OR3:.*]] = or i1 %[[COND_OR1]], %[[COND_OR2]]
+; CHECK-NEXT:    %[[COND_XOR:.*]] = xor i1 %cond5, %[[VAR_COND]]
+; CHECK-NEXT:    %[[COND_AND:.*]] = and i1 %cond6, %[[VAR_COND]]
+; CHECK-NEXT:    %[[COND_OR4:.*]] = or i1 %[[COND_XOR]], %[[COND_AND]]
+; CHECK-NEXT:    %[[COND_OR5:.*]] = or i1 %[[COND_OR3]], %[[COND_OR4]]
+; CHECK-NEXT:    %[[COND_OR6:.*]] = or i1 %[[COND_OR5]], false
+; CHECK-NEXT:    br i1 %[[COND_OR6]], label %loop_exit, label %do_something
+
+do_something:
+  call void @some_func() noreturn nounwind
+  br label %loop_begin
+; CHECK:       do_something:
+; CHECK-NEXT:    call
+; CHECK-NEXT:    br label %loop_begin
+
+loop_exit:
+  ret i32 0
+; CHECK:       loop_exit.split:
+; CHECK-NEXT:    ret
+}
+
+define i32 @test_partial_condition_unswitch_with_lcssa_phi1(i32* %var, i1 %cond, i32 %x) {
+; CHECK-LABEL: @test_partial_condition_unswitch_with_lcssa_phi1(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 %cond, label %entry.split, label %loop_exit.split
+;
+; CHECK:       entry.split:
+; CHECK-NEXT:    br label %loop_begin
+
+loop_begin:
+  %var_val = load i32, i32* %var
+  %var_cond = trunc i32 %var_val to i1
+  %cond_and = and i1 %var_cond, %cond
+  br i1 %cond_and, label %do_something, label %loop_exit
+; CHECK:       loop_begin:
+; CHECK-NEXT:    %[[VAR:.*]] = load i32
+; CHECK-NEXT:    %[[VAR_COND:.*]] = trunc i32 %[[VAR]] to i1
+; CHECK-NEXT:    %[[COND_AND:.*]] = and i1 %[[VAR_COND]], true
+; CHECK-NEXT:    br i1 %[[COND_AND]], label %do_something, label %loop_exit
+
+do_something:
+  call void @some_func() noreturn nounwind
+  br label %loop_begin
+; CHECK:       do_something:
+; CHECK-NEXT:    call
+; CHECK-NEXT:    br label %loop_begin
+
+loop_exit:
+  %x.lcssa = phi i32 [ %x, %loop_begin ]
+  ret i32 %x.lcssa
+; CHECK:       loop_exit:
+; CHECK-NEXT:    %[[LCSSA:.*]] = phi i32 [ %x, %loop_begin ]
+; CHECK-NEXT:    br label %loop_exit.split
+;
+; CHECK:       loop_exit.split:
+; CHECK-NEXT:    %[[LCSSA_SPLIT:.*]] = phi i32 [ %x, %entry ], [ %[[LCSSA]], %loop_exit ]
+; CHECK-NEXT:    ret i32 %[[LCSSA_SPLIT]]
+}
+
+define i32 @test_partial_condition_unswitch_with_lcssa_phi2(i32* %var, i1 %cond, i32 %x, i32 %y) {
+; CHECK-LABEL: @test_partial_condition_unswitch_with_lcssa_phi2(
+entry:
+  br label %loop_begin
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 %cond, label %entry.split, label %loop_exit.split
+;
+; CHECK:       entry.split:
+; CHECK-NEXT:    br label %loop_begin
+
+loop_begin:
+  %var_val = load i32, i32* %var
+  %var_cond = trunc i32 %var_val to i1
+  %cond_and = and i1 %var_cond, %cond
+  br i1 %cond_and, label %do_something, label %loop_exit
+; CHECK:       loop_begin:
+; CHECK-NEXT:    %[[VAR:.*]] = load i32
+; CHECK-NEXT:    %[[VAR_COND:.*]] = trunc i32 %[[VAR]] to i1
+; CHECK-NEXT:    %[[COND_AND:.*]] = and i1 %[[VAR_COND]], true
+; CHECK-NEXT:    br i1 %[[COND_AND]], label %do_something, label %loop_exit
+
+do_something:
+  call void @some_func() noreturn nounwind
+  br i1 %var_cond, label %loop_begin, label %loop_exit
+; CHECK:       do_something:
+; CHECK-NEXT:    call
+; CHECK-NEXT:    br i1 %[[VAR_COND]], label %loop_begin, label %loop_exit
+
+loop_exit:
+  %xy.lcssa = phi i32 [ %x, %loop_begin ], [ %y, %do_something ]
+  ret i32 %xy.lcssa
+; CHECK:       loop_exit:
+; CHECK-NEXT:    %[[LCSSA:.*]] = phi i32 [ %x, %loop_begin ], [ %y, %do_something ]
+; CHECK-NEXT:    br label %loop_exit.split
+;
+; CHECK:       loop_exit.split:
+; CHECK-NEXT:    %[[LCSSA_SPLIT:.*]] = phi i32 [ %x, %entry ], [ %[[LCSSA]], %loop_exit ]
+; CHECK-NEXT:    ret i32 %[[LCSSA_SPLIT]]
+}