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