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-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This pass lowers callbrs in LLVM IR in order to to assist SelectionDAG's 10 // codegen. 11 // 12 // In particular, this pass assists in inserting register copies for the output 13 // 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 in 15 // the IR representation, the value cannot be copied to the appropriate virtual 16 // registers prior to jumping to an indirect label, since the jump occurs 17 // within the user-provided assembly blob. 18 // 19 // Instead, those copies must occur separately at the beginning of each 20 // indirect target. That requires that we create a separate SSA definition in 21 // each of them (via llvm.callbr.landingpad), and may require splitting 22 // critical edges so we have a location to place the intrinsic. Finally, we 23 // remap users of the original callbr output SSA value to instead point to the 24 // appropriate llvm.callbr.landingpad value. 25 // 26 // Ideally, this could be done inside SelectionDAG, or in the 27 // 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 considering 30 // 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 53 using namespace llvm; 54 55 #define DEBUG_TYPE "callbr-prepare" 56 57 static bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT); 58 static bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs, 59 DominatorTree &DT); 60 static void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic, 61 SSAUpdater &SSAUpdate); 62 static SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn); 63 64 namespace { 65 66 class CallBrPrepare : public FunctionPass { 67 public: 68 CallBrPrepare() : FunctionPass(ID) {} 69 void getAnalysisUsage(AnalysisUsage &AU) const override; 70 bool runOnFunction(Function &Fn) override; 71 static char ID; 72 }; 73 74 } // end anonymous namespace 75 76 PreservedAnalyses CallBrPreparePass::run(Function &Fn, 77 FunctionAnalysisManager &FAM) { 78 bool Changed = false; 79 SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(Fn); 80 81 if (CBRs.empty()) 82 return PreservedAnalyses::all(); 83 84 auto &DT = FAM.getResult<DominatorTreeAnalysis>(Fn); 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 96 char CallBrPrepare::ID = 0; 97 INITIALIZE_PASS_BEGIN(CallBrPrepare, "callbrprepare", "Prepare callbr", false, 98 false) 99 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass) 100 INITIALIZE_PASS_END(CallBrPrepare, "callbrprepare", "Prepare callbr", false, 101 false) 102 103 FunctionPass *llvm::createCallBrPass() { return new CallBrPrepare(); } 104 105 void CallBrPrepare::getAnalysisUsage(AnalysisUsage &AU) const { 106 AU.addPreserved<DominatorTreeWrapperPass>(); 107 } 108 109 SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn) { 110 SmallVector<CallBrInst *, 2> CBRs; 111 for (BasicBlock &BB : Fn) 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 118 bool 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 need 126 // to split the default destination if it's duplicated between an indirect 127 // destination... 128 // %1 = callbr ... to label %x [label %x] 129 // ...hence starting at 1 and checking against successor 0 (aka the default 130 // 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 140 bool 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 167 static 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 NDEBUG 173 static 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 #endif 184 185 void 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 NDEBUG 198 PrintDebugDomInfo(DT, *U, LandingPad, /*IsDefaultDest*/ false); 199 PrintDebugDomInfo(DT, *U, DefaultDest, /*IsDefaultDest*/ true); 200 #endif 201 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, replace 208 // 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 222 bool CallBrPrepare::runOnFunction(Function &Fn) { 223 bool Changed = false; 224 SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(Fn); 225 226 if (CBRs.empty()) 227 return Changed; 228 229 // It's highly likely that most programs do not contain CallBrInsts. Follow a 230 // similar pattern from SafeStackLegacyPass::runOnFunction to reuse previous 231 // domtree analysis if available, otherwise compute it lazily. This avoids 232 // forcing Dominator Tree Construction at -O0 for programs that likely do not 233 // contain CallBrInsts. It does pessimize programs with callbr at higher 234 // optimization levels, as the DominatorTree created here is not reused by 235 // 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(Fn); 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