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/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:
  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
  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)
  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
 105 void CallBrPrepare::getAnalysisUsage(AnalysisUsage &AU) const {
 106   AU.addPreserved<DominatorTreeWrapperPass>();
 107 }
 108
 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
 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
 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
 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
 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
 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
 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 }