290 lines · 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-exception6//7//===----------------------------------------------------------------------===//8//9// For each natural loop with multiple exit blocks, this pass creates a new10// block N such that all exiting blocks now branch to N, and then control flow11// is redistributed to all the original exit blocks.12//13// Limitation: This assumes that all terminators in the CFG are direct branches14// (the "br" instruction). The presence of any other control flow15// such as indirectbr or switch will cause an assert.16// The callbr terminator is supported by creating intermediate17// target blocks that unconditionally branch to the original target18// blocks. These intermediate target blocks can then be redirected19// through the ControlFlowHub as usual.20//21//===----------------------------------------------------------------------===//22 23#include "llvm/Transforms/Utils/UnifyLoopExits.h"24#include "llvm/ADT/MapVector.h"25#include "llvm/Analysis/DomTreeUpdater.h"26#include "llvm/Analysis/LoopInfo.h"27#include "llvm/IR/Constants.h"28#include "llvm/IR/Dominators.h"29#include "llvm/InitializePasses.h"30#include "llvm/Support/CommandLine.h"31#include "llvm/Transforms/Utils.h"32#include "llvm/Transforms/Utils/BasicBlockUtils.h"33#include "llvm/Transforms/Utils/ControlFlowUtils.h"34 35#define DEBUG_TYPE "unify-loop-exits"36 37using namespace llvm;38 39static cl::opt<unsigned> MaxBooleansInControlFlowHub(40 "max-booleans-in-control-flow-hub", cl::init(32), cl::Hidden,41 cl::desc("Set the maximum number of outgoing blocks for using a boolean "42 "value to record the exiting block in the ControlFlowHub."));43 44namespace {45struct UnifyLoopExitsLegacyPass : public FunctionPass {46 static char ID;47 UnifyLoopExitsLegacyPass() : FunctionPass(ID) {48 initializeUnifyLoopExitsLegacyPassPass(*PassRegistry::getPassRegistry());49 }50 51 void getAnalysisUsage(AnalysisUsage &AU) const override {52 AU.addRequired<LoopInfoWrapperPass>();53 AU.addRequired<DominatorTreeWrapperPass>();54 AU.addPreserved<LoopInfoWrapperPass>();55 AU.addPreserved<DominatorTreeWrapperPass>();56 }57 58 bool runOnFunction(Function &F) override;59};60} // namespace61 62char UnifyLoopExitsLegacyPass::ID = 0;63 64FunctionPass *llvm::createUnifyLoopExitsPass() {65 return new UnifyLoopExitsLegacyPass();66}67 68INITIALIZE_PASS_BEGIN(UnifyLoopExitsLegacyPass, "unify-loop-exits",69 "Fixup each natural loop to have a single exit block",70 false /* Only looks at CFG */, false /* Analysis Pass */)71INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)72INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)73INITIALIZE_PASS_END(UnifyLoopExitsLegacyPass, "unify-loop-exits",74 "Fixup each natural loop to have a single exit block",75 false /* Only looks at CFG */, false /* Analysis Pass */)76 77// The current transform introduces new control flow paths which may break the78// SSA requirement that every def must dominate all its uses. For example,79// consider a value D defined inside the loop that is used by some instruction80// U outside the loop. It follows that D dominates U, since the original81// program has valid SSA form. After merging the exits, all paths from D to U82// now flow through the unified exit block. In addition, there may be other83// paths that do not pass through D, but now reach the unified exit84// block. Thus, D no longer dominates U.85//86// Restore the dominance by creating a phi for each such D at the new unified87// loop exit. But when doing this, ignore any uses U that are in the new unified88// loop exit, since those were introduced specially when the block was created.89//90// The use of SSAUpdater seems like overkill for this operation. The location91// for creating the new PHI is well-known, and also the set of incoming blocks92// to the new PHI.93static void restoreSSA(const DominatorTree &DT, const Loop *L,94 SmallVectorImpl<BasicBlock *> &Incoming,95 BasicBlock *LoopExitBlock) {96 using InstVector = SmallVector<Instruction *, 8>;97 using IIMap = MapVector<Instruction *, InstVector>;98 IIMap ExternalUsers;99 for (auto *BB : L->blocks()) {100 for (auto &I : *BB) {101 for (auto &U : I.uses()) {102 auto UserInst = cast<Instruction>(U.getUser());103 auto UserBlock = UserInst->getParent();104 if (UserBlock == LoopExitBlock)105 continue;106 if (L->contains(UserBlock))107 continue;108 LLVM_DEBUG(dbgs() << "added ext use for " << I.getName() << "("109 << BB->getName() << ")"110 << ": " << UserInst->getName() << "("111 << UserBlock->getName() << ")"112 << "\n");113 ExternalUsers[&I].push_back(UserInst);114 }115 }116 }117 118 for (const auto &II : ExternalUsers) {119 // For each Def used outside the loop, create NewPhi in120 // LoopExitBlock. NewPhi receives Def only along exiting blocks that121 // dominate it, while the remaining values are undefined since those paths122 // didn't exist in the original CFG.123 auto Def = II.first;124 LLVM_DEBUG(dbgs() << "externally used: " << Def->getName() << "\n");125 auto NewPhi =126 PHINode::Create(Def->getType(), Incoming.size(),127 Def->getName() + ".moved", LoopExitBlock->begin());128 for (auto *In : Incoming) {129 LLVM_DEBUG(dbgs() << "predecessor " << In->getName() << ": ");130 if (Def->getParent() == In || DT.dominates(Def, In)) {131 LLVM_DEBUG(dbgs() << "dominated\n");132 NewPhi->addIncoming(Def, In);133 } else {134 LLVM_DEBUG(dbgs() << "not dominated\n");135 NewPhi->addIncoming(PoisonValue::get(Def->getType()), In);136 }137 }138 139 LLVM_DEBUG(dbgs() << "external users:");140 for (auto *U : II.second) {141 LLVM_DEBUG(dbgs() << " " << U->getName());142 U->replaceUsesOfWith(Def, NewPhi);143 }144 LLVM_DEBUG(dbgs() << "\n");145 }146}147 148static bool unifyLoopExits(DominatorTree &DT, LoopInfo &LI, Loop *L) {149 // To unify the loop exits, we need a list of the exiting blocks as150 // well as exit blocks. The functions for locating these lists both151 // traverse the entire loop body. It is more efficient to first152 // locate the exiting blocks and then examine their successors to153 // locate the exit blocks.154 SmallVector<BasicBlock *, 8> ExitingBlocks;155 L->getExitingBlocks(ExitingBlocks);156 157 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);158 SmallVector<BasicBlock *, 8> CallBrTargetBlocksToFix;159 // Redirect exiting edges through a control flow hub.160 ControlFlowHub CHub;161 bool Changed = false;162 163 for (unsigned I = 0; I < ExitingBlocks.size(); ++I) {164 BasicBlock *BB = ExitingBlocks[I];165 if (BranchInst *Branch = dyn_cast<BranchInst>(BB->getTerminator())) {166 BasicBlock *Succ0 = Branch->getSuccessor(0);167 Succ0 = L->contains(Succ0) ? nullptr : Succ0;168 169 BasicBlock *Succ1 =170 Branch->isUnconditional() ? nullptr : Branch->getSuccessor(1);171 Succ1 = L->contains(Succ1) ? nullptr : Succ1;172 CHub.addBranch(BB, Succ0, Succ1);173 174 LLVM_DEBUG(dbgs() << "Added extiting branch: " << printBasicBlock(BB)175 << " -> " << printBasicBlock(Succ0)176 << (Succ0 && Succ1 ? " " : "") << printBasicBlock(Succ1)177 << '\n');178 } else if (CallBrInst *CallBr = dyn_cast<CallBrInst>(BB->getTerminator())) {179 for (unsigned J = 0; J < CallBr->getNumSuccessors(); ++J) {180 BasicBlock *Succ = CallBr->getSuccessor(J);181 if (L->contains(Succ))182 continue;183 bool UpdatedLI = false;184 BasicBlock *NewSucc =185 SplitCallBrEdge(BB, Succ, J, &DTU, nullptr, &LI, &UpdatedLI);186 // SplitCallBrEdge modifies the CFG because it creates an intermediate187 // block. So we need to set the changed flag no matter what the188 // ControlFlowHub is going to do later.189 Changed = true;190 // Even if CallBr and Succ do not have a common parent loop, we need to191 // add the new target block to the parent loop of the current loop.192 if (!UpdatedLI)193 CallBrTargetBlocksToFix.push_back(NewSucc);194 // ExitingBlocks is later used to restore SSA, so we need to make sure195 // that the blocks used for phi nodes in the guard blocks match the196 // predecessors of the guard blocks, which, in the case of callbr, are197 // the new intermediate target blocks instead of the callbr blocks198 // themselves.199 ExitingBlocks[I] = NewSucc;200 CHub.addBranch(NewSucc, Succ);201 LLVM_DEBUG(dbgs() << "Added exiting branch: "202 << printBasicBlock(NewSucc) << " -> "203 << printBasicBlock(Succ) << '\n');204 }205 } else {206 llvm_unreachable("unsupported block terminator");207 }208 }209 210 SmallVector<BasicBlock *, 8> GuardBlocks;211 BasicBlock *LoopExitBlock;212 bool ChangedCFG;213 std::tie(LoopExitBlock, ChangedCFG) = CHub.finalize(214 &DTU, GuardBlocks, "loop.exit", MaxBooleansInControlFlowHub.getValue());215 ChangedCFG |= Changed;216 if (!ChangedCFG)217 return false;218 219 restoreSSA(DT, L, ExitingBlocks, LoopExitBlock);220 221#if defined(EXPENSIVE_CHECKS)222 assert(DT.verify(DominatorTree::VerificationLevel::Full));223#else224 assert(DT.verify(DominatorTree::VerificationLevel::Fast));225#endif // EXPENSIVE_CHECKS226 L->verifyLoop();227 228 // The guard blocks were created outside the loop, so they need to become229 // members of the parent loop.230 // Same goes for the callbr target blocks. Although we try to add them to the231 // smallest common parent loop of the callbr block and the corresponding232 // original target block, there might not have been such a loop, in which case233 // the newly created callbr target blocks are not part of any loop. For nested234 // loops, this might result in them leading to a loop with multiple entry235 // points.236 if (auto *ParentLoop = L->getParentLoop()) {237 for (auto *G : GuardBlocks) {238 ParentLoop->addBasicBlockToLoop(G, LI);239 }240 for (auto *C : CallBrTargetBlocksToFix) {241 ParentLoop->addBasicBlockToLoop(C, LI);242 }243 ParentLoop->verifyLoop();244 }245 246#if defined(EXPENSIVE_CHECKS)247 LI.verify(DT);248#endif // EXPENSIVE_CHECKS249 250 return true;251}252 253static bool runImpl(LoopInfo &LI, DominatorTree &DT) {254 255 bool Changed = false;256 auto Loops = LI.getLoopsInPreorder();257 for (auto *L : Loops) {258 LLVM_DEBUG(dbgs() << "Processing loop:\n"; L->print(dbgs()));259 Changed |= unifyLoopExits(DT, LI, L);260 }261 return Changed;262}263 264bool UnifyLoopExitsLegacyPass::runOnFunction(Function &F) {265 LLVM_DEBUG(dbgs() << "===== Unifying loop exits in function " << F.getName()266 << "\n");267 auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();268 auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();269 270 return runImpl(LI, DT);271}272 273namespace llvm {274 275PreservedAnalyses UnifyLoopExitsPass::run(Function &F,276 FunctionAnalysisManager &AM) {277 LLVM_DEBUG(dbgs() << "===== Unifying loop exits in function " << F.getName()278 << "\n");279 auto &LI = AM.getResult<LoopAnalysis>(F);280 auto &DT = AM.getResult<DominatorTreeAnalysis>(F);281 282 if (!runImpl(LI, DT))283 return PreservedAnalyses::all();284 PreservedAnalyses PA;285 PA.preserve<LoopAnalysis>();286 PA.preserve<DominatorTreeAnalysis>();287 return PA;288}289} // namespace llvm290