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 "callbrprepare" 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, DEBUG_TYPE, "Prepare callbr", false, false) 98 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass) 99 INITIALIZE_PASS_END(CallBrPrepare, DEBUG_TYPE, "Prepare callbr", false, false) 100 101 FunctionPass *llvm::createCallBrPass() { return new CallBrPrepare(); } 102 103 void CallBrPrepare::getAnalysisUsage(AnalysisUsage &AU) const { 104 AU.addPreserved<DominatorTreeWrapperPass>(); 105 } 106 107 SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn) { 108 SmallVector<CallBrInst *, 2> CBRs; 109 for (BasicBlock &BB : Fn) 110 if (auto *CBR = dyn_cast<CallBrInst>(BB.getTerminator())) 111 if (!CBR->getType()->isVoidTy() && !CBR->use_empty()) 112 CBRs.push_back(CBR); 113 return CBRs; 114 } 115 116 bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT) { 117 bool Changed = false; 118 CriticalEdgeSplittingOptions Options(&DT); 119 Options.setMergeIdenticalEdges(); 120 121 // The indirect destination might be duplicated between another parameter... 122 // %0 = callbr ... [label %x, label %x] 123 // ...hence MergeIdenticalEdges and AllowIndentical edges, but we don't need 124 // to split the default destination if it's duplicated between an indirect 125 // destination... 126 // %1 = callbr ... to label %x [label %x] 127 // ...hence starting at 1 and checking against successor 0 (aka the default 128 // destination). 129 for (CallBrInst *CBR : CBRs) 130 for (unsigned i = 1, e = CBR->getNumSuccessors(); i != e; ++i) 131 if (CBR->getSuccessor(i) == CBR->getSuccessor(0) || 132 isCriticalEdge(CBR, i, /*AllowIdenticalEdges*/ true)) 133 if (SplitKnownCriticalEdge(CBR, i, Options)) 134 Changed = true; 135 return Changed; 136 } 137 138 bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT) { 139 bool Changed = false; 140 SmallPtrSet<const BasicBlock *, 4> Visited; 141 IRBuilder<> Builder(CBRs[0]->getContext()); 142 for (CallBrInst *CBR : CBRs) { 143 if (!CBR->getNumIndirectDests()) 144 continue; 145 146 SSAUpdater SSAUpdate; 147 SSAUpdate.Initialize(CBR->getType(), CBR->getName()); 148 SSAUpdate.AddAvailableValue(CBR->getParent(), CBR); 149 SSAUpdate.AddAvailableValue(CBR->getDefaultDest(), CBR); 150 151 for (BasicBlock *IndDest : CBR->getIndirectDests()) { 152 if (!Visited.insert(IndDest).second) 153 continue; 154 Builder.SetInsertPoint(&*IndDest->begin()); 155 CallInst *Intrinsic = Builder.CreateIntrinsic( 156 CBR->getType(), Intrinsic::callbr_landingpad, {CBR}); 157 SSAUpdate.AddAvailableValue(IndDest, Intrinsic); 158 UpdateSSA(DT, CBR, Intrinsic, SSAUpdate); 159 Changed = true; 160 } 161 } 162 return Changed; 163 } 164 165 static bool IsInSameBasicBlock(const Use &U, const BasicBlock *BB) { 166 const auto *I = dyn_cast<Instruction>(U.getUser()); 167 return I && I->getParent() == BB; 168 } 169 170 #ifndef NDEBUG 171 static void PrintDebugDomInfo(const DominatorTree &DT, const Use &U, 172 const BasicBlock *BB, bool IsDefaultDest) { 173 if (!isa<Instruction>(U.getUser())) 174 return; 175 LLVM_DEBUG(dbgs() << "Use: " << *U.getUser() << ", in block " 176 << cast<Instruction>(U.getUser())->getParent()->getName() 177 << ", is " << (DT.dominates(BB, U) ? "" : "NOT ") 178 << "dominated by " << BB->getName() << " (" 179 << (IsDefaultDest ? "in" : "") << "direct)\n"); 180 } 181 #endif 182 183 void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic, 184 SSAUpdater &SSAUpdate) { 185 186 SmallPtrSet<Use *, 4> Visited; 187 BasicBlock *DefaultDest = CBR->getDefaultDest(); 188 BasicBlock *LandingPad = Intrinsic->getParent(); 189 190 SmallVector<Use *, 4> Uses(make_pointer_range(CBR->uses())); 191 for (Use *U : Uses) { 192 if (!Visited.insert(U).second) 193 continue; 194 195 #ifndef NDEBUG 196 PrintDebugDomInfo(DT, *U, LandingPad, /*IsDefaultDest*/ false); 197 PrintDebugDomInfo(DT, *U, DefaultDest, /*IsDefaultDest*/ true); 198 #endif 199 200 // Don't rewrite the use in the newly inserted intrinsic. 201 if (const auto *II = dyn_cast<IntrinsicInst>(U->getUser())) 202 if (II->getIntrinsicID() == Intrinsic::callbr_landingpad) 203 continue; 204 205 // If the Use is in the same BasicBlock as the Intrinsic call, replace 206 // the Use with the value of the Intrinsic call. 207 if (IsInSameBasicBlock(*U, LandingPad)) { 208 U->set(Intrinsic); 209 continue; 210 } 211 212 // If the Use is dominated by the default dest, do not touch it. 213 if (DT.dominates(DefaultDest, *U)) 214 continue; 215 216 SSAUpdate.RewriteUse(*U); 217 } 218 } 219 220 bool CallBrPrepare::runOnFunction(Function &Fn) { 221 bool Changed = false; 222 SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(Fn); 223 224 if (CBRs.empty()) 225 return Changed; 226 227 // It's highly likely that most programs do not contain CallBrInsts. Follow a 228 // similar pattern from SafeStackLegacyPass::runOnFunction to reuse previous 229 // domtree analysis if available, otherwise compute it lazily. This avoids 230 // forcing Dominator Tree Construction at -O0 for programs that likely do not 231 // contain CallBrInsts. It does pessimize programs with callbr at higher 232 // optimization levels, as the DominatorTree created here is not reused by 233 // subsequent passes. 234 DominatorTree *DT; 235 std::optional<DominatorTree> LazilyComputedDomTree; 236 if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>()) 237 DT = &DTWP->getDomTree(); 238 else { 239 LazilyComputedDomTree.emplace(Fn); 240 DT = &*LazilyComputedDomTree; 241 } 242 243 if (SplitCriticalEdges(CBRs, *DT)) 244 Changed = true; 245 246 if (InsertIntrinsicCalls(CBRs, *DT)) 247 Changed = true; 248 249 return Changed; 250 } 251