[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

Why not s/provided vector/DeadInsts?

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

76

Nit: wrap?

322

Stray line?

333

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

398

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

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

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

Wasn't done?

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

249

"just walk"

269

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

446

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.

452

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

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

269

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

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

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

269

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

446

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

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 ↗

(On Diff #152128)

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 ↗

(On Diff #152128)

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 ↗

(On Diff #152128)

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 ↗

(On Diff #152128)

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 ↗

(On Diff #152128)

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 ↗

(On Diff #152128)

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/include/llvm/Analysis/Utils/Local.h

bool wouldInstructionBeTriviallyDead(Instruction *I,
                                      const TargetLibraryInfo *TLI = nullptr);
 
 /// If the specified value is a trivially dead instruction, delete it.
 /// If that makes any of its operands trivially dead, delete them too,
 /// recursively. Return true if any instructions were deleted.
 bool RecursivelyDeleteTriviallyDeadInstructions(Value *V,
                                         const TargetLibraryInfo *TLI = nullptr);
 
+/// Delete all of the instructions in the provided vector, and all other
+/// instructions that deleting these in turn causes to be trivially dead.
+///
+/// The initial instructions in the provided vector must all have empty use
+/// lists and satisfy `isInstructionTriviallyDead`.
+///
+/// The provided `SmallVector` will be used as scratch storage for this routine
+/// and will be empty afterward.
+void RecursivelyDeleteTriviallyDeadInstructions(
+    SmallVectorImpl<Instruction *> &DeadInsts,
+    const TargetLibraryInfo *TLI = nullptr);
+
 /// If the specified value is an effectively dead PHI node, due to being a
 /// def-use chain of single-use nodes that either forms a cycle or is terminated
 /// by a trivially dead instruction, delete it. If that makes any of its
 /// operands trivially dead, delete them too, recursively. Return true if a
 /// change was made.
 bool RecursivelyDeleteDeadPHINode(PHINode *PN,
                                   const TargetLibraryInfo *TLI = nullptr);
 

llvm/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,
+/// 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;
+
+  // 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 replaceLoopInvariantUsesAndSimplify(Loop &L,
+                                                ArrayRef<Value *> Invariants,
                                                 Constant &Replacement) {
-  assert(!isa<Constant>(LIC) && "Why are we unswitching on a constant?");
+  // Keep a list of instructions to simplify after replacing uses of invariants
+  // within the loop.
+  SmallPtrSet<Instruction *, 4> SimplifiedSet;
+  SmallVector<Instruction *, 4> SimplificationWorklist;
+
+  for (Value *Invariant : Invariants) {
+    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;
+
+      // Track that this is something we may need to simplify.
+      if (SimplifiedSet.insert(UserI).second)
+        SimplificationWorklist.push_back(UserI);
     }
   }
 
+  Module &M = *L.getHeader()->getParent()->getParent();
+  auto &DL = M.getDataLayout();
+
+  SmallVector<Instruction *, 4> DeadInsts;
+  while (!SimplificationWorklist.empty()) {
+    Instruction &I = *SimplificationWorklist.pop_back_val();
+    // First check if this is already dead and skip it.
+    if (I.use_empty()) {
+      if (isInstructionTriviallyDead(&I))
+        DeadInsts.push_back(&I);
+      continue;
+    }
+
+    // Try to do very basic simplifications of these instructions.
+    Value *SimpleV = SimplifyInstruction(&I, SimplifyQuery(DL, &I));
+    if (!SimpleV)
+      continue;
+
+    // Manually do RAUW so that we can intercept users within the loop and
+    // update them.
+    for (auto UI = I.use_begin(), UE = I.use_end(); UI != UE;) {
+      // Grab the use and walk past it so we can clobber it in the use list.
+      Use *U = &*UI++;
+
+      // Replace the use with the simplified value.
+      *U = SimpleV;
+
+      // If this is an instruction within the loop, recurse through it. The
+      // reason we restrict this to in-loop uses is that we don't want to
+      // simplify the LCSSA phi node away. We could solve this, but 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.
+      Instruction *UserI = dyn_cast<Instruction>(U->getUser());
+      if (UserI && L.contains(UserI) && SimplifiedSet.insert(UserI).second)
+        SimplificationWorklist.push_back(UserI);
+    }
+
+    // We expect the instruction to no longer have uses, so check if it can be
+    // deleted.
+    assert(I.use_empty() && "Didn't rewrite a use after simplification?");
+    if (isInstructionTriviallyDead(&I))
+      DeadInsts.push_back(&I);
+  }
+
+  RecursivelyDeleteTriviallyDeadInstructions(DeadInsts);
+}
+
 /// 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) {
     auto *PN = dyn_cast<PHINode>(&I);
     if (!PN)
       // No more PHIs to check.

 /// 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) {
   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);
       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!");
   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);
+  } else {
+    // 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;
+  // ConstantInt *CondVal = ConstantInt::getTrue(BI.getContext());
+  // ConstantInt *Replacement = ConstantInt::getFalse(BI.getContext());
   int LoopExitSuccIdx = 0;
   auto *LoopExitBB = BI.getSuccessor(0);
   if (L.contains(LoopExitBB)) {
-    std::swap(CondVal, Replacement);
+    ExitDirection = false;
+    // std::swap(CondVal, Replacement);
     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;
 
-  DEBUG(dbgs() << "    unswitching trivial branch when: " << CondVal
-               << " == " << LoopCond << "\n");
+  // For partial unswitching of a 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;
+    }
+  }
+
+  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();
+    for (Value *Invariant :
+         make_range(std::next(Invariants.begin()), Invariants.end()))
+      if (LoopExitSuccIdx == 0)
+        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);
 
   // 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);
+  replaceLoopInvariantUsesAndSimplify(L, Invariants, *Replacement);
 
   ++NumTrivial;
   ++NumBranches;
   return true;
 }
 
 /// Unswitch a trivial switch if the condition is loop invariant.
 ///
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/lib/Transforms/Utils/Local.cpp

 llvm::RecursivelyDeleteTriviallyDeadInstructions(Value *V,
                                                  const TargetLibraryInfo *TLI) {
   Instruction *I = dyn_cast<Instruction>(V);
   if (!I || !I->use_empty() || !isInstructionTriviallyDead(I, TLI))
     return false;
 
   SmallVector<Instruction*, 16> DeadInsts;
   DeadInsts.push_back(I);
+  RecursivelyDeleteTriviallyDeadInstructions(DeadInsts, TLI);
 
-  do {
-    I = DeadInsts.pop_back_val();
-    salvageDebugInfo(*I);
+  return true;
+}
+
+void llvm::RecursivelyDeleteTriviallyDeadInstructions(
+    SmallVectorImpl<Instruction *> &DeadInsts, const TargetLibraryInfo *TLI) {
+  // Process the dead instruction list until empty.
+  while (!DeadInsts.empty()) {
+    Instruction &I = *DeadInsts.pop_back_val();
+    salvageDebugInfo(I);
 
     // Null out all of the instruction's operands to see if any operand becomes
     // dead as we go.
-    for (unsigned i = 0, e = I->getNumOperands(); i != e; ++i) {
-      Value *OpV = I->getOperand(i);
-      I->setOperand(i, nullptr);
+    for (Use &OpU : I.operands()) {
+      Value *OpV = OpU.get();
+      OpU = nullptr;
 
-      if (!OpV->use_empty()) continue;
+      if (!OpV->use_empty())
+        continue;
 
       // If the operand is an instruction that became dead as we nulled out the
       // operand, and if it is 'trivially' dead, delete it in a future loop
       // iteration.
       if (Instruction *OpI = dyn_cast<Instruction>(OpV))
         if (isInstructionTriviallyDead(OpI, TLI))
           DeadInsts.push_back(OpI);
     }
 
-    I->eraseFromParent();
-  } while (!DeadInsts.empty());
-
-  return true;
+    I.eraseFromParent();
+  }
 }
 
 /// areAllUsesEqual - Check whether the uses of a value are all the same.
 /// This is similar to Instruction::hasOneUse() except this will also return
 /// true when there are no uses or multiple uses that all refer to the same
 /// value.
 static bool areAllUsesEqual(Instruction *I) {
   Value::user_iterator UI = I->user_begin();

llvm/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:    br i1 %[[COND2]], 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/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:    %[[VAL:.*]] = load i32
+; CHECK-NEXT:    %[[VAL_COND:.*]] = trunc i32 %[[VAL]] to i1
+; CHECK-NEXT:    br i1 %[[VAL_COND]], 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) {
+; 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_or4 = or i1 %cond_or3, %cond4
+  br i1 %cond_or4, label %loop_exit, label %do_something
+; CHECK:       loop_begin:
+; CHECK-NEXT:    %[[VAL:.*]] = load i32
+; CHECK-NEXT:    %[[VAL_COND:.*]] = trunc i32 %[[VAL]] to i1
+; CHECK-NEXT:    br i1 %[[VAL_COND]], 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
+}
 No newline at end of file