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
53using namespace llvm;
54
55#define DEBUG_TYPE "callbrprepare"
56
57static bool SplitCriticalEdges(ArrayRef<CallBrInst *> CBRs, DominatorTree &DT);
58static bool InsertIntrinsicCalls(ArrayRef<CallBrInst *> CBRs,
59                                 DominatorTree &DT);
60static void UpdateSSA(DominatorTree &DT, CallBrInst *CBR, CallInst *Intrinsic,
61                      SSAUpdater &SSAUpdate);
62static SmallVector<CallBrInst *, 2> FindCallBrs(Function &Fn);
63
64namespace {
65
66class CallBrPrepare : public FunctionPass {
67public:
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
76PreservedAnalyses 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
96char CallBrPrepare::ID = 0;
97INITIALIZE_PASS_BEGIN(CallBrPrepare, DEBUG_TYPE, "Prepare callbr", false, false)
98INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
99INITIALIZE_PASS_END(CallBrPrepare, DEBUG_TYPE, "Prepare callbr", false, false)
100
101FunctionPass *llvm::createCallBrPass() { return new CallBrPrepare(); }
102
103void CallBrPrepare::getAnalysisUsage(AnalysisUsage &AU) const {
104  AU.addPreserved<DominatorTreeWrapperPass>();
105}
106
107SmallVector<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
116bool 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
138bool 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
165static 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
171static 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
183void 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
220bool 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