llvm/lib/CodeGen/CallBrPrepare.cpp

   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 }