[SimpleLoopUnswitch] Make all checks before actual non-trivial unswitch

[mkazantsev]

Summary

We should be able to make all relevant checks before we actually start the non-trivial
unswitching, so that we could guarantee that once we have started the transform,
it will always succeed.

[chandlerc]

FWIW, after my suggestion I would not call this "NFC" as it should change the cost function of this routine at least....

lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

2368–2376 ↗

(On Diff #171260)

Can we hoist this even fiurther up? Maybe w/ the irreducible CFG?

We shouldn't even compute and rank the cost of the candidates when this is the case.

The reason this was inside the unswitch (ages ago) was because we were re-using teh exit blocks list and I didn't want to create it twice, but your solution of passing it in seems way better.

[chandlerc]

Awesome improvement, thanks!

LGTM. Minor tweak or follow-up suggested below.

lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

1799 ↗

(On Diff #171271)

Maybe make this come before the dom tree? I'd generally prefer roughly grouped params... but SCEV is already a bit out-of-order, so also fine to clean up the grouping of parameters in a follow-up, whatever is easiest. I think the best ordering here would be: loop, terminator, invarianst, exit blocks, callback, analyses-used.

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

do {
     for (DomTreeNode *ChildN : *N) {
       assert(Visited.insert(ChildN).second &&
              "Cannot visit a node twice when walking a tree!");
       DomWorklist.push_back(ChildN);
     }
   } while (!DomWorklist.empty());
 }
 
-static bool unswitchNontrivialInvariants(
+static void unswitchNontrivialInvariants(
     Loop &L, Instruction &TI, ArrayRef<Value *> Invariants,
-    DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
-    function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB,
+    SmallVectorImpl<BasicBlock *> &ExitBlocks, DominatorTree &DT, LoopInfo &LI,
+    AssumptionCache &AC, function_ref<void(bool, ArrayRef<Loop *>)> UnswitchCB,
     ScalarEvolution *SE) {
   auto *ParentBB = TI.getParent();
   BranchInst *BI = dyn_cast<BranchInst>(&TI);
   SwitchInst *SI = BI ? nullptr : cast<SwitchInst>(&TI);
 
   // We can only unswitch switches, conditional branches with an invariant
   // condition, or combining invariant conditions with an instruction.
   assert((SI || BI->isConditional()) &&
assert(!UnswitchedSuccBBs.count(RetainedSuccBB) &&
          "Should not unswitch the same successor we are retaining!");
 
   // The branch should be in this exact loop. Any inner loop's invariant branch
   // should be handled by unswitching that inner loop. The caller of this
   // routine should filter out any candidates that remain (but were skipped for
   // whatever reason).
   assert(LI.getLoopFor(ParentBB) == &L && "Branch in an inner loop!");
 
-  SmallVector<BasicBlock *, 4> ExitBlocks;
-  L.getUniqueExitBlocks(ExitBlocks);
-
-  // We cannot unswitch if exit blocks contain a cleanuppad instruction as we
-  // don't know how to split those exit blocks.
-  // FIXME: We should teach SplitBlock to handle this and remove this
-  // restriction.
-  for (auto *ExitBB : ExitBlocks)
-    if (isa<CleanupPadInst>(ExitBB->getFirstNonPHI()))
-      return false;
-
   // Compute the parent loop now before we start hacking on things.
   Loop *ParentL = L.getParentLoop();
 
   // Compute the outer-most loop containing one of our exit blocks. This is the
   // furthest up our loopnest which can be mutated, which we will use below to
   // update things.
   Loop *OuterExitL = &L;
   for (auto *ExitBB : ExitBlocks) {
#endif
   // and check whether the original loop survived.
   SmallVector<Loop *, 4> SibLoops;
   for (Loop *UpdatedL : llvm::concat<Loop *>(NonChildClonedLoops, HoistedLoops))
     if (UpdatedL->getParentLoop() == ParentL)
       SibLoops.push_back(UpdatedL);
   UnswitchCB(IsStillLoop, SibLoops);
 
   ++NumBranches;
-  return true;
 }
 
 /// Recursively compute the cost of a dominator subtree based on the per-block
 /// cost map provided.
 ///
 /// The recursive computation is memozied into the provided DT-indexed cost map
 /// to allow querying it for most nodes in the domtree without it becoming
 /// quadratic.
unswitchBestCondition(Loop &L, DominatorTree &DT, LoopInfo &LI,
   // loops "out of thin air". If we ever discover important use cases for doing
   // this, we can add support to loop unswitch, but it is a lot of complexity
   // for what seems little or no real world benefit.
   LoopBlocksRPO RPOT(&L);
   RPOT.perform(&LI);
   if (containsIrreducibleCFG<const BasicBlock *>(RPOT, LI))
     return false;
 
+  SmallVector<BasicBlock *, 4> ExitBlocks;
+  L.getUniqueExitBlocks(ExitBlocks);
+
+  // We cannot unswitch if exit blocks contain a cleanuppad instruction as we
+  // don't know how to split those exit blocks.
+  // FIXME: We should teach SplitBlock to handle this and remove this
+  // restriction.
+  for (auto *ExitBB : ExitBlocks)
+    if (isa<CleanupPadInst>(ExitBB->getFirstNonPHI())) {
+      dbgs() << "Cannot unswitch because of cleanuppad in exit block\n";
+      return false;
+    }
+
   LLVM_DEBUG(
       dbgs() << "Considering " << UnswitchCandidates.size()
              << " non-trivial loop invariant conditions for unswitching.\n");
 
   // Given that unswitching these terminators will require duplicating parts of
   // the loop, so we need to be able to model that cost. Compute the ephemeral
   // values and set up a data structure to hold per-BB costs. We cache each
   // block's cost so that we don't recompute this when considering different
unswitchBestCondition(Loop &L, DominatorTree &DT, LoopInfo &LI,
   }
 
   if (BestUnswitchCost >= UnswitchThreshold) {
     LLVM_DEBUG(dbgs() << "Cannot unswitch, lowest cost found: "
                       << BestUnswitchCost << "\n");
     return false;
   }
 
-  LLVM_DEBUG(dbgs() << "  Trying to unswitch non-trivial (cost = "
+  LLVM_DEBUG(dbgs() << "  Unswitching non-trivial (cost = "
                     << BestUnswitchCost << ") terminator: " << *BestUnswitchTI
                     << "\n");
-  return unswitchNontrivialInvariants(
-      L, *BestUnswitchTI, BestUnswitchInvariants, DT, LI, AC, UnswitchCB, SE);
+  unswitchNontrivialInvariants(L, *BestUnswitchTI, BestUnswitchInvariants,
+                               ExitBlocks, DT, LI, AC, UnswitchCB, SE);
+  return true;
 }
 
 /// Unswitch control flow predicated on loop invariant conditions.
 ///
 /// This first hoists all branches or switches which are trivial (IE, do not
 /// require duplicating any part of the loop) out of the loop body. It then
 /// looks at other loop invariant control flows and tries to unswitch those as
 /// well by cloning the loop if the result is small enough.