253 lines · cpp
1//===-- CallBrPrepare - Prepare callbr for code generation ----------------===//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// This pass lowers callbrs in LLVM IR in order to to assist SelectionDAG's10// codegen.11//12// In particular, this pass assists in inserting register copies for the output13// values of a callbr along the edges leading to the indirect target blocks.14// Though the output SSA value is defined by the callbr instruction itself in15// the IR representation, the value cannot be copied to the appropriate virtual16// registers prior to jumping to an indirect label, since the jump occurs17// within the user-provided assembly blob.18//19// Instead, those copies must occur separately at the beginning of each20// indirect target. That requires that we create a separate SSA definition in21// each of them (via llvm.callbr.landingpad), and may require splitting22// critical edges so we have a location to place the intrinsic. Finally, we23// remap users of the original callbr output SSA value to instead point to the24// appropriate llvm.callbr.landingpad value.25//26// Ideally, this could be done inside SelectionDAG, or in the27// MachineInstruction representation, without the use of an IR-level intrinsic.28// But, within the current framework, it’s simpler to implement as an IR pass.29// (If support for callbr in GlobalISel is implemented, it’s worth considering30// whether this is still required.)31//32//===----------------------------------------------------------------------===//33 34#include "llvm/CodeGen/CallBrPrepare.h"35#include "llvm/ADT/ArrayRef.h"36#include "llvm/ADT/SmallPtrSet.h"37#include "llvm/ADT/SmallVector.h"38#include "llvm/ADT/iterator.h"39#include "llvm/Analysis/CFG.h"40#include "llvm/CodeGen/Passes.h"41#include "llvm/IR/BasicBlock.h"42#include "llvm/IR/Dominators.h"43#include "llvm/IR/Function.h"44#include "llvm/IR/IRBuilder.h"45#include "llvm/IR/Instructions.h"46#include "llvm/IR/IntrinsicInst.h"47#include "llvm/IR/Intrinsics.h"48#include "llvm/InitializePasses.h"49#include "llvm/Pass.h"50#include "llvm/Transforms/Utils/BasicBlockUtils.h"51#include "llvm/Transforms/Utils/SSAUpdater.h"52 53using namespace llvm;54 55#define DEBUG_TYPE "callbr-prepare"56 57static bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT);58static bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs,59 DominatorTree &DT);60static void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic,61 SSAUpdater &SSAUpdate);62static SmallVector<CallBrInst *, 2> FindCallBrs(Function &F);63 64namespace {65 66class CallBrPrepare : public FunctionPass {67public:68 CallBrPrepare() : FunctionPass(ID) {}69 void getAnalysisUsage(AnalysisUsage &AU) const override;70 bool runOnFunction(Function &F) override;71 static char ID;72};73 74} // end anonymous namespace75 76PreservedAnalyses CallBrPreparePass::run(Function &F,77 FunctionAnalysisManager &FAM) {78 bool Changed = false;79 SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(F);80 81 if (CBRs.empty())82 return PreservedAnalyses::all();83 84 auto &DT = FAM.getResult<DominatorTreeAnalysis>(F);85 86 Changed |= SplitCriticalEdges(CBRs, DT);87 Changed |= InsertIntrinsicCalls(CBRs, DT);88 89 if (!Changed)90 return PreservedAnalyses::all();91 PreservedAnalyses PA;92 PA.preserve<DominatorTreeAnalysis>();93 return PA;94}95 96char CallBrPrepare::ID = 0;97INITIALIZE_PASS_BEGIN(CallBrPrepare, "callbrprepare", "Prepare callbr", false,98 false)99INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)100INITIALIZE_PASS_END(CallBrPrepare, "callbrprepare", "Prepare callbr", false,101 false)102 103FunctionPass *llvm::createCallBrPass() { return new CallBrPrepare(); }104 105void CallBrPrepare::getAnalysisUsage(AnalysisUsage &AU) const {106 AU.addPreserved<DominatorTreeWrapperPass>();107}108 109SmallVector<CallBrInst *, 2> FindCallBrs(Function &F) {110 SmallVector<CallBrInst *, 2> CBRs;111 for (BasicBlock &BB : F)112 if (auto *CBR = dyn_cast<CallBrInst>(BB.getTerminator()))113 if (!CBR->getType()->isVoidTy() && !CBR->use_empty())114 CBRs.push_back(CBR);115 return CBRs;116}117 118bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT) {119 bool Changed = false;120 CriticalEdgeSplittingOptions Options(&DT);121 Options.setMergeIdenticalEdges();122 123 // The indirect destination might be duplicated between another parameter...124 // %0 = callbr ... [label %x, label %x]125 // ...hence MergeIdenticalEdges and AllowIndentical edges, but we don't need126 // to split the default destination if it's duplicated between an indirect127 // destination...128 // %1 = callbr ... to label %x [label %x]129 // ...hence starting at 1 and checking against successor 0 (aka the default130 // destination).131 for (CallBrInst *CBR : CBRs)132 for (unsigned i = 1, e = CBR->getNumSuccessors(); i != e; ++i)133 if (CBR->getSuccessor(i) == CBR->getSuccessor(0) ||134 isCriticalEdge(CBR, i, /*AllowIdenticalEdges*/ true))135 if (SplitKnownCriticalEdge(CBR, i, Options))136 Changed = true;137 return Changed;138}139 140bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT) {141 bool Changed = false;142 SmallPtrSet<const BasicBlock *, 4> Visited;143 IRBuilder<> Builder(CBRs[0]->getContext());144 for (CallBrInst *CBR : CBRs) {145 if (!CBR->getNumIndirectDests())146 continue;147 148 SSAUpdater SSAUpdate;149 SSAUpdate.Initialize(CBR->getType(), CBR->getName());150 SSAUpdate.AddAvailableValue(CBR->getParent(), CBR);151 SSAUpdate.AddAvailableValue(CBR->getDefaultDest(), CBR);152 153 for (BasicBlock *IndDest : CBR->getIndirectDests()) {154 if (!Visited.insert(IndDest).second)155 continue;156 Builder.SetInsertPoint(&*IndDest->begin());157 CallInst *Intrinsic = Builder.CreateIntrinsic(158 CBR->getType(), Intrinsic::callbr_landingpad, {CBR});159 SSAUpdate.AddAvailableValue(IndDest, Intrinsic);160 UpdateSSA(DT, CBR, Intrinsic, SSAUpdate);161 Changed = true;162 }163 }164 return Changed;165}166 167static bool IsInSameBasicBlock(const Use &U, const BasicBlock *BB) {168 const auto *I = dyn_cast<Instruction>(U.getUser());169 return I && I->getParent() == BB;170}171 172#ifndef NDEBUG173static void PrintDebugDomInfo(const DominatorTree &DT, const Use &U,174 const BasicBlock *BB, bool IsDefaultDest) {175 if (!isa<Instruction>(U.getUser()))176 return;177 LLVM_DEBUG(dbgs() << "Use: " << *U.getUser() << ", in block "178 << cast<Instruction>(U.getUser())->getParent()->getName()179 << ", is " << (DT.dominates(BB, U) ? "" : "NOT ")180 << "dominated by " << BB->getName() << " ("181 << (IsDefaultDest ? "in" : "") << "direct)\n");182}183#endif184 185void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic,186 SSAUpdater &SSAUpdate) {187 188 SmallPtrSet<Use *, 4> Visited;189 BasicBlock *DefaultDest = CBR->getDefaultDest();190 BasicBlock *LandingPad = Intrinsic->getParent();191 192 SmallVector<Use *, 4> Uses(make_pointer_range(CBR->uses()));193 for (Use *U : Uses) {194 if (!Visited.insert(U).second)195 continue;196 197#ifndef NDEBUG198 PrintDebugDomInfo(DT, *U, LandingPad, /*IsDefaultDest*/ false);199 PrintDebugDomInfo(DT, *U, DefaultDest, /*IsDefaultDest*/ true);200#endif201 202 // Don't rewrite the use in the newly inserted intrinsic.203 if (const auto *II = dyn_cast<IntrinsicInst>(U->getUser()))204 if (II->getIntrinsicID() == Intrinsic::callbr_landingpad)205 continue;206 207 // If the Use is in the same BasicBlock as the Intrinsic call, replace208 // the Use with the value of the Intrinsic call.209 if (IsInSameBasicBlock(*U, LandingPad)) {210 U->set(Intrinsic);211 continue;212 }213 214 // If the Use is dominated by the default dest, do not touch it.215 if (DT.dominates(DefaultDest, *U))216 continue;217 218 SSAUpdate.RewriteUse(*U);219 }220}221 222bool CallBrPrepare::runOnFunction(Function &F) {223 bool Changed = false;224 SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(F);225 226 if (CBRs.empty())227 return Changed;228 229 // It's highly likely that most programs do not contain CallBrInsts. Follow a230 // similar pattern from SafeStackLegacyPass::runOnFunction to reuse previous231 // domtree analysis if available, otherwise compute it lazily. This avoids232 // forcing Dominator Tree Construction at -O0 for programs that likely do not233 // contain CallBrInsts. It does pessimize programs with callbr at higher234 // optimization levels, as the DominatorTree created here is not reused by235 // subsequent passes.236 DominatorTree *DT;237 std::optional<DominatorTree> LazilyComputedDomTree;238 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())239 DT = &DTWP->getDomTree();240 else {241 LazilyComputedDomTree.emplace(F);242 DT = &*LazilyComputedDomTree;243 }244 245 if (SplitCriticalEdges(CBRs, *DT))246 Changed = true;247 248 if (InsertIntrinsicCalls(CBRs, *DT))249 Changed = true;250 251 return Changed;252}253