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/ADT/ArrayRef.h" |
35 | #include "llvm/ADT/SmallPtrSet.h" |
36 | #include "llvm/ADT/SmallVector.h" |
37 | #include "llvm/ADT/iterator.h" |
38 | #include "llvm/Analysis/CFG.h" |
39 | #include "llvm/CodeGen/Passes.h" |
40 | #include "llvm/IR/BasicBlock.h" |
41 | #include "llvm/IR/Dominators.h" |
42 | #include "llvm/IR/Function.h" |
43 | #include "llvm/IR/IRBuilder.h" |
44 | #include "llvm/IR/Instructions.h" |
45 | #include "llvm/IR/IntrinsicInst.h" |
46 | #include "llvm/IR/Intrinsics.h" |
47 | #include "llvm/InitializePasses.h" |
48 | #include "llvm/Pass.h" |
49 | #include "llvm/Transforms/Utils/BasicBlockUtils.h" |
50 | #include "llvm/Transforms/Utils/SSAUpdater.h" |
51 | |
52 | using namespace llvm; |
53 | |
54 | #define DEBUG_TYPE "callbrprepare" |
55 | |
56 | namespace { |
57 | |
58 | class CallBrPrepare : public FunctionPass { |
59 | bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT); |
60 | bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs, |
61 | DominatorTree &DT) const; |
62 | void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic, |
63 | SSAUpdater &SSAUpdate) const; |
64 | |
65 | public: |
66 | CallBrPrepare() : FunctionPass(ID) {} |
67 | void getAnalysisUsage(AnalysisUsage &AU) const override; |
68 | bool runOnFunction(Function &Fn) override; |
69 | static char ID; |
70 | }; |
71 | |
72 | } // end anonymous namespace |
73 | |
74 | char CallBrPrepare::ID = 0; |
75 | INITIALIZE_PASS_BEGIN(CallBrPrepare, DEBUG_TYPE, "Prepare callbr" , false, false) |
76 | INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass) |
77 | INITIALIZE_PASS_END(CallBrPrepare, DEBUG_TYPE, "Prepare callbr" , false, false) |
78 | |
79 | FunctionPass *llvm::createCallBrPass() { return new CallBrPrepare(); } |
80 | |
81 | void CallBrPrepare::getAnalysisUsage(AnalysisUsage &AU) const { |
82 | AU.addPreserved<DominatorTreeWrapperPass>(); |
83 | } |
84 | |
85 | static SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn) { |
86 | SmallVector<CallBrInst *, 2> CBRs; |
87 | for (BasicBlock &BB : Fn) |
88 | if (auto *CBR = dyn_cast<CallBrInst>(BB.getTerminator())) |
89 | if (!CBR->getType()->isVoidTy() && !CBR->use_empty()) |
90 | CBRs.push_back(CBR); |
91 | return CBRs; |
92 | } |
93 | |
94 | bool CallBrPrepare::SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, |
95 | DominatorTree &DT) { |
96 | bool Changed = false; |
97 | CriticalEdgeSplittingOptions Options(&DT); |
98 | Options.setMergeIdenticalEdges(); |
99 | |
100 | // The indirect destination might be duplicated between another parameter... |
101 | // %0 = callbr ... [label %x, label %x] |
102 | // ...hence MergeIdenticalEdges and AllowIndentical edges, but we don't need |
103 | // to split the default destination if it's duplicated between an indirect |
104 | // destination... |
105 | // %1 = callbr ... to label %x [label %x] |
106 | // ...hence starting at 1 and checking against successor 0 (aka the default |
107 | // destination). |
108 | for (CallBrInst *CBR : CBRs) |
109 | for (unsigned i = 1, e = CBR->getNumSuccessors(); i != e; ++i) |
110 | if (CBR->getSuccessor(i) == CBR->getSuccessor(0) || |
111 | isCriticalEdge(CBR, i, /*AllowIdenticalEdges*/ true)) |
112 | if (SplitKnownCriticalEdge(CBR, i, Options)) |
113 | Changed = true; |
114 | return Changed; |
115 | } |
116 | |
117 | bool CallBrPrepare::InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs, |
118 | DominatorTree &DT) const { |
119 | bool Changed = false; |
120 | SmallPtrSet<const BasicBlock *, 4> Visited; |
121 | IRBuilder<> Builder(CBRs[0]->getContext()); |
122 | for (CallBrInst *CBR : CBRs) { |
123 | if (!CBR->getNumIndirectDests()) |
124 | continue; |
125 | |
126 | SSAUpdater SSAUpdate; |
127 | SSAUpdate.Initialize(CBR->getType(), CBR->getName()); |
128 | SSAUpdate.AddAvailableValue(CBR->getParent(), CBR); |
129 | SSAUpdate.AddAvailableValue(CBR->getDefaultDest(), CBR); |
130 | |
131 | for (BasicBlock *IndDest : CBR->getIndirectDests()) { |
132 | if (!Visited.insert(IndDest).second) |
133 | continue; |
134 | Builder.SetInsertPoint(&*IndDest->begin()); |
135 | CallInst *Intrinsic = Builder.CreateIntrinsic( |
136 | CBR->getType(), Intrinsic::callbr_landingpad, {CBR}); |
137 | SSAUpdate.AddAvailableValue(IndDest, Intrinsic); |
138 | UpdateSSA(DT, CBR, Intrinsic, SSAUpdate); |
139 | Changed = true; |
140 | } |
141 | } |
142 | return Changed; |
143 | } |
144 | |
145 | static bool IsInSameBasicBlock(const Use &U, const BasicBlock *BB) { |
146 | const auto *I = dyn_cast<Instruction>(U.getUser()); |
147 | return I && I->getParent() == BB; |
148 | } |
149 | |
150 | #ifndef NDEBUG |
151 | static void PrintDebugDomInfo(const DominatorTree &DT, const Use &U, |
152 | const BasicBlock *BB, bool IsDefaultDest) { |
153 | if (!isa<Instruction>(U.getUser())) |
154 | return; |
155 | LLVM_DEBUG(dbgs() << "Use: " << *U.getUser() << ", in block " |
156 | << cast<Instruction>(U.getUser())->getParent()->getName() |
157 | << ", is " << (DT.dominates(BB, U) ? "" : "NOT " ) |
158 | << "dominated by " << BB->getName() << " (" |
159 | << (IsDefaultDest ? "in" : "" ) << "direct)\n" ); |
160 | } |
161 | #endif |
162 | |
163 | void CallBrPrepare::UpdateSSA(DominatorTree &DT, CallBrInst *CBR, |
164 | CallInst *Intrinsic, |
165 | SSAUpdater &SSAUpdate) const { |
166 | |
167 | SmallPtrSet<Use *, 4> Visited; |
168 | BasicBlock *DefaultDest = CBR->getDefaultDest(); |
169 | BasicBlock *LandingPad = Intrinsic->getParent(); |
170 | |
171 | SmallVector<Use *, 4> Uses(make_pointer_range(CBR->uses())); |
172 | for (Use *U : Uses) { |
173 | if (!Visited.insert(U).second) |
174 | continue; |
175 | |
176 | #ifndef NDEBUG |
177 | PrintDebugDomInfo(DT, *U, LandingPad, /*IsDefaultDest*/ false); |
178 | PrintDebugDomInfo(DT, *U, DefaultDest, /*IsDefaultDest*/ true); |
179 | #endif |
180 | |
181 | // Don't rewrite the use in the newly inserted intrinsic. |
182 | if (const auto *II = dyn_cast<IntrinsicInst>(U->getUser())) |
183 | if (II->getIntrinsicID() == Intrinsic::callbr_landingpad) |
184 | continue; |
185 | |
186 | // If the Use is in the same BasicBlock as the Intrinsic call, replace |
187 | // the Use with the value of the Intrinsic call. |
188 | if (IsInSameBasicBlock(*U, LandingPad)) { |
189 | U->set(Intrinsic); |
190 | continue; |
191 | } |
192 | |
193 | // If the Use is dominated by the default dest, do not touch it. |
194 | if (DT.dominates(DefaultDest, *U)) |
195 | continue; |
196 | |
197 | SSAUpdate.RewriteUse(*U); |
198 | } |
199 | } |
200 | |
201 | bool CallBrPrepare::runOnFunction(Function &Fn) { |
202 | bool Changed = false; |
203 | SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(Fn); |
204 | |
205 | if (CBRs.empty()) |
206 | return Changed; |
207 | |
208 | // It's highly likely that most programs do not contain CallBrInsts. Follow a |
209 | // similar pattern from SafeStackLegacyPass::runOnFunction to reuse previous |
210 | // domtree analysis if available, otherwise compute it lazily. This avoids |
211 | // forcing Dominator Tree Construction at -O0 for programs that likely do not |
212 | // contain CallBrInsts. It does pessimize programs with callbr at higher |
213 | // optimization levels, as the DominatorTree created here is not reused by |
214 | // subsequent passes. |
215 | DominatorTree *DT; |
216 | std::optional<DominatorTree> LazilyComputedDomTree; |
217 | if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>()) |
218 | DT = &DTWP->getDomTree(); |
219 | else { |
220 | LazilyComputedDomTree.emplace(Fn); |
221 | DT = &*LazilyComputedDomTree; |
222 | } |
223 | |
224 | if (SplitCriticalEdges(CBRs, *DT)) |
225 | Changed = true; |
226 | |
227 | if (InsertIntrinsicCalls(CBRs, *DT)) |
228 | Changed = true; |
229 | |
230 | return Changed; |
231 | } |
232 | |