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:
CallBrPrepare()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
run(Function & Fn,FunctionAnalysisManager & FAM)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)
INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)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
getAnalysisUsage(AnalysisUsage & AU) const105 void CallBrPrepare::getAnalysisUsage(AnalysisUsage &AU) const {
106 AU.addPreserved<DominatorTreeWrapperPass>();
107 }
108
FindCallBrs(Function & Fn)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
SplitCriticalEdges(ArrayRef<CallBrInst * > CBRs,DominatorTree & DT)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
InsertIntrinsicCalls(ArrayRef<CallBrInst * > CBRs,DominatorTree & DT)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
IsInSameBasicBlock(const Use & U,const BasicBlock * BB)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
PrintDebugDomInfo(const DominatorTree & DT,const Use & U,const BasicBlock * BB,bool IsDefaultDest)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
UpdateSSA(DominatorTree & DT,CallBrInst * CBR,CallInst * Intrinsic,SSAUpdater & SSAUpdate)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
runOnFunction(Function & Fn)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