xref: /llvm-project/llvm/lib/CodeGen/CallBrPrepare.cpp (revision 094190c2f52900d9d5d26dba9522e70030cc00a1)
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/Analysis/CFG.h"
38 #include "llvm/CodeGen/Passes.h"
39 #include "llvm/IR/BasicBlock.h"
40 #include "llvm/IR/Dominators.h"
41 #include "llvm/IR/Function.h"
42 #include "llvm/IR/IRBuilder.h"
43 #include "llvm/IR/Instructions.h"
44 #include "llvm/IR/IntrinsicInst.h"
45 #include "llvm/IR/Intrinsics.h"
46 #include "llvm/InitializePasses.h"
47 #include "llvm/Pass.h"
48 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
49 
50 using namespace llvm;
51 
52 #define DEBUG_TYPE "callbrprepare"
53 
54 namespace {
55 
56 class CallBrPrepare : public FunctionPass {
57   bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT);
58 
59 public:
60   CallBrPrepare() : FunctionPass(ID) {}
61   void getAnalysisUsage(AnalysisUsage &AU) const override;
62   bool runOnFunction(Function &Fn) override;
63   static char ID;
64 };
65 
66 } // end anonymous namespace
67 
68 char CallBrPrepare::ID = 0;
69 INITIALIZE_PASS_BEGIN(CallBrPrepare, DEBUG_TYPE, "Prepare callbr", false, false)
70 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
71 INITIALIZE_PASS_END(CallBrPrepare, DEBUG_TYPE, "Prepare callbr", false, false)
72 
73 FunctionPass *llvm::createCallBrPass() { return new CallBrPrepare(); }
74 
75 void CallBrPrepare::getAnalysisUsage(AnalysisUsage &AU) const {
76   AU.addPreserved<DominatorTreeWrapperPass>();
77 }
78 
79 static SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn) {
80   SmallVector<CallBrInst *, 2> CBRs;
81   for (BasicBlock &BB : Fn)
82     if (auto *CBR = dyn_cast<CallBrInst>(BB.getTerminator()))
83       if (!CBR->getType()->isVoidTy() && !CBR->use_empty())
84         CBRs.push_back(CBR);
85   return CBRs;
86 }
87 
88 bool CallBrPrepare::SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs,
89                                        DominatorTree &DT) {
90   bool Changed = false;
91   CriticalEdgeSplittingOptions Options(&DT);
92   Options.setMergeIdenticalEdges();
93 
94   // The indirect destination might be duplicated between another parameter...
95   //   %0 = callbr ... [label %x, label %x]
96   // ...hence MergeIdenticalEdges and AllowIndentical edges, but we don't need
97   // to split the default destination if it's duplicated between an indirect
98   // destination...
99   //   %1 = callbr ... to label %x [label %x]
100   // ...hence starting at 1 and checking against successor 0 (aka the default
101   // destination).
102   for (CallBrInst *CBR : CBRs)
103     for (unsigned i = 1, e = CBR->getNumSuccessors(); i != e; ++i)
104       if (CBR->getSuccessor(i) == CBR->getSuccessor(0) ||
105           isCriticalEdge(CBR, i, /*AllowIdenticalEdges*/ true))
106         if (SplitKnownCriticalEdge(CBR, i, Options))
107           Changed = true;
108   return Changed;
109 }
110 
111 static bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs) {
112   bool Changed = false;
113   SmallPtrSet<const BasicBlock *, 4> Visited;
114   IRBuilder<> Builder(CBRs[0]->getContext());
115   for (CallBrInst *CBR : CBRs) {
116     for (BasicBlock *IndDest : CBR->getIndirectDests()) {
117       if (!Visited.insert(IndDest).second)
118         continue;
119       Builder.SetInsertPoint(&*IndDest->begin());
120       Builder.CreateIntrinsic(CBR->getType(), Intrinsic::callbr_landingpad,
121                               {CBR});
122       Changed = true;
123     }
124   }
125   return Changed;
126 }
127 
128 bool CallBrPrepare::runOnFunction(Function &Fn) {
129   bool Changed = false;
130   SmallVector<CallBrInst *, 2> CBRs = FindCallBrs(Fn);
131 
132   if (CBRs.empty())
133     return Changed;
134 
135   // It's highly likely that most programs do not contain CallBrInsts. Follow a
136   // similar pattern from SafeStackLegacyPass::runOnFunction to reuse previous
137   // domtree analysis if available, otherwise compute it lazily. This avoids
138   // forcing Dominator Tree Construction at -O0 for programs that likely do not
139   // contain CallBrInsts. It does pessimize programs with callbr at higher
140   // optimization levels, as the DominatorTree created here is not reused by
141   // subsequent passes.
142   DominatorTree *DT;
143   std::optional<DominatorTree> LazilyComputedDomTree;
144   if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
145     DT = &DTWP->getDomTree();
146   else {
147     LazilyComputedDomTree.emplace(Fn);
148     DT = &*LazilyComputedDomTree;
149   }
150 
151   if (SplitCriticalEdges(CBRs, *DT))
152     Changed = true;
153 
154   if (InsertIntrinsicCalls(CBRs))
155     Changed = true;
156 
157   return Changed;
158 }
159