llvm/lib/Transforms/Utils/UnifyLoopExits.cpp

   1 //===- UnifyLoopExits.cpp - Redirect exiting edges to one block -*- C++ -*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // For each natural loop with multiple exit blocks, this pass creates a new
  10 // block N such that all exiting blocks now branch to N, and then control flow
  11 // is redistributed to all the original exit blocks.
  12 //
  13 // Limitation: This assumes that all terminators in the CFG are direct branches
  14 //             (the "br" instruction). The presence of any other control flow
  15 //             such as indirectbr, switch or callbr will cause an assert.
  16 //
  17 //===----------------------------------------------------------------------===//
  18
  19 #include "llvm/Transforms/Utils/UnifyLoopExits.h"
  20 #include "llvm/ADT/MapVector.h"
  21 #include "llvm/Analysis/DomTreeUpdater.h"
  22 #include "llvm/Analysis/LoopInfo.h"
  23 #include "llvm/IR/Constants.h"
  24 #include "llvm/IR/Dominators.h"
  25 #include "llvm/InitializePasses.h"
  26 #include "llvm/Support/CommandLine.h"
  27 #include "llvm/Transforms/Utils.h"
  28 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
  29 #include "llvm/Transforms/Utils/ControlFlowUtils.h"
  30
  31 #define DEBUG_TYPE "unify-loop-exits"
  32
  33 using namespace llvm;
  34
  35 static cl::opt<unsigned> MaxBooleansInControlFlowHub(
  36     "max-booleans-in-control-flow-hub", cl::init(32), cl::Hidden,
  37     cl::desc("Set the maximum number of outgoing blocks for using a boolean "
  38              "value to record the exiting block in the ControlFlowHub."));
  39
  40 namespace {
  41 struct UnifyLoopExitsLegacyPass : public FunctionPass {
  42   static char ID;
  43   UnifyLoopExitsLegacyPass() : FunctionPass(ID) {
  44     initializeUnifyLoopExitsLegacyPassPass(*PassRegistry::getPassRegistry());
  45   }
  46
  47   void getAnalysisUsage(AnalysisUsage &AU) const override {
  48     AU.addRequired<LoopInfoWrapperPass>();
  49     AU.addRequired<DominatorTreeWrapperPass>();
  50     AU.addPreserved<LoopInfoWrapperPass>();
  51     AU.addPreserved<DominatorTreeWrapperPass>();
  52   }
  53
  54   bool runOnFunction(Function &F) override;
  55 };
  56 } // namespace
  57
  58 char UnifyLoopExitsLegacyPass::ID = 0;
  59
  60 FunctionPass *llvm::createUnifyLoopExitsPass() {
  61   return new UnifyLoopExitsLegacyPass();
  62 }
  63
  64 INITIALIZE_PASS_BEGIN(UnifyLoopExitsLegacyPass, "unify-loop-exits",
  65                       "Fixup each natural loop to have a single exit block",
  66                       false /* Only looks at CFG */, false /* Analysis Pass */)
  67 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
  68 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
  69 INITIALIZE_PASS_END(UnifyLoopExitsLegacyPass, "unify-loop-exits",
  70                     "Fixup each natural loop to have a single exit block",
  71                     false /* Only looks at CFG */, false /* Analysis Pass */)
  72
  73 // The current transform introduces new control flow paths which may break the
  74 // SSA requirement that every def must dominate all its uses. For example,
  75 // consider a value D defined inside the loop that is used by some instruction
  76 // U outside the loop. It follows that D dominates U, since the original
  77 // program has valid SSA form. After merging the exits, all paths from D to U
  78 // now flow through the unified exit block. In addition, there may be other
  79 // paths that do not pass through D, but now reach the unified exit
  80 // block. Thus, D no longer dominates U.
  81 //
  82 // Restore the dominance by creating a phi for each such D at the new unified
  83 // loop exit. But when doing this, ignore any uses U that are in the new unified
  84 // loop exit, since those were introduced specially when the block was created.
  85 //
  86 // The use of SSAUpdater seems like overkill for this operation. The location
  87 // for creating the new PHI is well-known, and also the set of incoming blocks
  88 // to the new PHI.
  89 static void restoreSSA(const DominatorTree &DT, const Loop *L,
  90                        SmallVectorImpl<BasicBlock *> &Incoming,
  91                        BasicBlock *LoopExitBlock) {
  92   using InstVector = SmallVector<Instruction *, 8>;
  93   using IIMap = MapVector<Instruction *, InstVector>;
  94   IIMap ExternalUsers;
  95   for (auto *BB : L->blocks()) {
  96     for (auto &I : *BB) {
  97       for (auto &U : I.uses()) {
  98         auto UserInst = cast<Instruction>(U.getUser());
  99         auto UserBlock = UserInst->getParent();
 100         if (UserBlock == LoopExitBlock)
 101           continue;
 102         if (L->contains(UserBlock))
 103           continue;
 104         LLVM_DEBUG(dbgs() << "added ext use for " << I.getName() << "("
 105                           << BB->getName() << ")"
 106                           << ": " << UserInst->getName() << "("
 107                           << UserBlock->getName() << ")"
 108                           << "\n");
 109         ExternalUsers[&I].push_back(UserInst);
 110       }
 111     }
 112   }
 113
 114   for (const auto &II : ExternalUsers) {
 115     // For each Def used outside the loop, create NewPhi in
 116     // LoopExitBlock. NewPhi receives Def only along exiting blocks that
 117     // dominate it, while the remaining values are undefined since those paths
 118     // didn't exist in the original CFG.
 119     auto Def = II.first;
 120     LLVM_DEBUG(dbgs() << "externally used: " << Def->getName() << "\n");
 121     auto NewPhi =
 122         PHINode::Create(Def->getType(), Incoming.size(),
 123                         Def->getName() + ".moved", LoopExitBlock->begin());
 124     for (auto *In : Incoming) {
 125       LLVM_DEBUG(dbgs() << "predecessor " << In->getName() << ": ");
 126       if (Def->getParent() == In || DT.dominates(Def, In)) {
 127         LLVM_DEBUG(dbgs() << "dominated\n");
 128         NewPhi->addIncoming(Def, In);
 129       } else {
 130         LLVM_DEBUG(dbgs() << "not dominated\n");
 131         NewPhi->addIncoming(PoisonValue::get(Def->getType()), In);
 132       }
 133     }
 134
 135     LLVM_DEBUG(dbgs() << "external users:");
 136     for (auto *U : II.second) {
 137       LLVM_DEBUG(dbgs() << " " << U->getName());
 138       U->replaceUsesOfWith(Def, NewPhi);
 139     }
 140     LLVM_DEBUG(dbgs() << "\n");
 141   }
 142 }
 143
 144 static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {
 145   // To unify the loop exits, we need a list of the exiting blocks as
 146   // well as exit blocks. The functions for locating these lists both
 147   // traverse the entire loop body. It is more efficient to first
 148   // locate the exiting blocks and then examine their successors to
 149   // locate the exit blocks.
 150   SmallVector<BasicBlock *, 8> ExitingBlocks;
 151   L->getExitingBlocks(ExitingBlocks);
 152
 153   // Redirect exiting edges through a control flow hub.
 154   ControlFlowHub CHub;
 155   for (auto *BB : ExitingBlocks) {
 156     auto *Branch = cast<BranchInst>(BB->getTerminator());
 157     BasicBlock *Succ0 = Branch->getSuccessor(0);
 158     Succ0 = L->contains(Succ0) ? nullptr : Succ0;
 159
 160     BasicBlock *Succ1 =
 161         Branch->isUnconditional() ? nullptr : Branch->getSuccessor(1);
 162     Succ1 = L->contains(Succ1) ? nullptr : Succ1;
 163     CHub.addBranch(BB, Succ0, Succ1);
 164
 165     LLVM_DEBUG(dbgs() << "Added exiting branch: " << BB->getName() << " -> {"
 166                       << (Succ0 ? Succ0->getName() : "<none>") << ", "
 167                       << (Succ1 ? Succ1->getName() : "<none>") << "}\n");
 168   }
 169
 170   SmallVector<BasicBlock *, 8> GuardBlocks;
 171   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
 172   BasicBlock *LoopExitBlock = CHub.finalize(
 173       &DTU, GuardBlocks, "loop.exit", MaxBooleansInControlFlowHub.getValue());
 174
 175   restoreSSA(DT, L, ExitingBlocks, LoopExitBlock);
 176
 177 #if defined(EXPENSIVE_CHECKS)
 178   assert(DT.verify(DominatorTree::VerificationLevel::Full));
 179 #else
 180   assert(DT.verify(DominatorTree::VerificationLevel::Fast));
 181 #endif // EXPENSIVE_CHECKS
 182   L->verifyLoop();
 183
 184   // The guard blocks were created outside the loop, so they need to become
 185   // members of the parent loop.
 186   if (auto ParentLoop = L->getParentLoop()) {
 187     for (auto *G : GuardBlocks) {
 188       ParentLoop->addBasicBlockToLoop(G, LI);
 189     }
 190     ParentLoop->verifyLoop();
 191   }
 192
 193 #if defined(EXPENSIVE_CHECKS)
 194   LI.verify(DT);
 195 #endif // EXPENSIVE_CHECKS
 196
 197   return true;
 198 }
 199
 200 static bool runImpl(LoopInfo &LI, DominatorTree &DT) {
 201
 202   bool Changed = false;
 203   auto Loops = LI.getLoopsInPreorder();
 204   for (auto *L : Loops) {
 205     LLVM_DEBUG(dbgs() << "Processing loop:\n"; L->print(dbgs()));
 206     Changed |= unifyLoopExits(DT, LI, L);
 207   }
 208   return Changed;
 209 }
 210
 211 bool UnifyLoopExitsLegacyPass::runOnFunction(Function &F) {
 212   LLVM_DEBUG(dbgs() << "===== Unifying loop exits in function " << F.getName()
 213                     << "\n");
 214   auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
 215   auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
 216
 217   assert(hasOnlySimpleTerminator(F) && "Unsupported block terminator.");
 218
 219   return runImpl(LI, DT);
 220 }
 221
 222 namespace llvm {
 223
 224 PreservedAnalyses UnifyLoopExitsPass::run(Function &F,
 225                                           FunctionAnalysisManager &AM) {
 226   LLVM_DEBUG(dbgs() << "===== Unifying loop exits in function " << F.getName()
 227                     << "\n");
 228   auto &LI = AM.getResult<LoopAnalysis>(F);
 229   auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
 230
 231   if (!runImpl(LI, DT))
 232     return PreservedAnalyses::all();
 233   PreservedAnalyses PA;
 234   PA.preserve<LoopAnalysis>();
 235   PA.preserve<DominatorTreeAnalysis>();
 236   return PA;
 237 }
 238 } // namespace llvm