[PGO] Make indirect call promotion a utility

This patch factors out the main code transformation utilities in the pgo-driven
indirect call promotion pass and places them in Transforms/Utils. The change is
intended to be a non-functional change, letting non-pgo-driven passes share a
common implementation with the existing pgo-driven pass.

The common utilities are used to conditionally promote indirect call sites to
direct call sites. They perform the underlying transformation, and do not
consider profile information. The pgo-specific details (e.g., the computation
of branch weight metadata) have been left in the indirect call promotion pass.

Differential Revision: https://reviews.llvm.org/D40658

llvm-svn: 319963
diff --git a/llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp b/llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp
index 8b9bbb4..49b8a67 100644
--- a/llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp
+++ b/llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp
@@ -47,6 +47,7 @@
 #include "llvm/Transforms/Instrumentation.h"
 #include "llvm/Transforms/PGOInstrumentation.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/CallPromotionUtils.h"
 #include <cassert>
 #include <cstdint>
 #include <memory>
@@ -214,49 +215,6 @@
 
 } // end anonymous namespace
 
-bool llvm::isLegalToPromote(Instruction *Inst, Function *F,
-                            const char **Reason) {
-  // Check the return type.
-  Type *CallRetType = Inst->getType();
-  if (!CallRetType->isVoidTy()) {
-    Type *FuncRetType = F->getReturnType();
-    if (FuncRetType != CallRetType &&
-        !CastInst::isBitCastable(FuncRetType, CallRetType)) {
-      if (Reason)
-        *Reason = "Return type mismatch";
-      return false;
-    }
-  }
-
-  // Check if the arguments are compatible with the parameters
-  FunctionType *DirectCalleeType = F->getFunctionType();
-  unsigned ParamNum = DirectCalleeType->getFunctionNumParams();
-  CallSite CS(Inst);
-  unsigned ArgNum = CS.arg_size();
-
-  if (ParamNum != ArgNum && !DirectCalleeType->isVarArg()) {
-    if (Reason)
-      *Reason = "The number of arguments mismatch";
-    return false;
-  }
-
-  for (unsigned I = 0; I < ParamNum; ++I) {
-    Type *PTy = DirectCalleeType->getFunctionParamType(I);
-    Type *ATy = CS.getArgument(I)->getType();
-    if (PTy == ATy)
-      continue;
-    if (!CastInst::castIsValid(Instruction::BitCast, CS.getArgument(I), PTy)) {
-      if (Reason)
-        *Reason = "Argument type mismatch";
-      return false;
-    }
-  }
-
-  DEBUG(dbgs() << " #" << NumOfPGOICallPromotion << " Promote the icall to "
-               << F->getName() << "\n");
-  return true;
-}
-
 // Indirect-call promotion heuristic. The direct targets are sorted based on
 // the count. Stop at the first target that is not promoted.
 std::vector<ICallPromotionFunc::PromotionCandidate>
@@ -317,7 +275,7 @@
     }
 
     const char *Reason = nullptr;
-    if (!isLegalToPromote(Inst, TargetFunction, &Reason)) {
+    if (!isLegalToPromote(CallSite(Inst), TargetFunction, &Reason)) {
       using namespace ore;
 
       ORE.emit([&]() {
@@ -335,23 +293,11 @@
   return Ret;
 }
 
-// Create a diamond structure for If_Then_Else. Also update the profile
-// count. Do the fix-up for the invoke instruction.
-static void createIfThenElse(Instruction *Inst, Function *DirectCallee,
-                             uint64_t Count, uint64_t TotalCount,
-                             BasicBlock **DirectCallBB,
-                             BasicBlock **IndirectCallBB,
-                             BasicBlock **MergeBB) {
-  CallSite CS(Inst);
-  Value *OrigCallee = CS.getCalledValue();
-
-  IRBuilder<> BBBuilder(Inst);
-  LLVMContext &Ctx = Inst->getContext();
-  Value *BCI1 =
-      BBBuilder.CreateBitCast(OrigCallee, Type::getInt8PtrTy(Ctx), "");
-  Value *BCI2 =
-      BBBuilder.CreateBitCast(DirectCallee, Type::getInt8PtrTy(Ctx), "");
-  Value *PtrCmp = BBBuilder.CreateICmpEQ(BCI1, BCI2, "");
+Instruction *llvm::pgo::promoteIndirectCall(Instruction *Inst,
+                                            Function *DirectCallee,
+                                            uint64_t Count, uint64_t TotalCount,
+                                            bool AttachProfToDirectCall,
+                                            OptimizationRemarkEmitter *ORE) {
 
   uint64_t ElseCount = TotalCount - Count;
   uint64_t MaxCount = (Count >= ElseCount ? Count : ElseCount);
@@ -359,231 +305,9 @@
   MDBuilder MDB(Inst->getContext());
   MDNode *BranchWeights = MDB.createBranchWeights(
       scaleBranchCount(Count, Scale), scaleBranchCount(ElseCount, Scale));
-  TerminatorInst *ThenTerm, *ElseTerm;
-  SplitBlockAndInsertIfThenElse(PtrCmp, Inst, &ThenTerm, &ElseTerm,
-                                BranchWeights);
-  *DirectCallBB = ThenTerm->getParent();
-  (*DirectCallBB)->setName("if.true.direct_targ");
-  *IndirectCallBB = ElseTerm->getParent();
-  (*IndirectCallBB)->setName("if.false.orig_indirect");
-  *MergeBB = Inst->getParent();
-  (*MergeBB)->setName("if.end.icp");
 
-  // Special handing of Invoke instructions.
-  InvokeInst *II = dyn_cast<InvokeInst>(Inst);
-  if (!II)
-    return;
-
-  // We don't need branch instructions for invoke.
-  ThenTerm->eraseFromParent();
-  ElseTerm->eraseFromParent();
-
-  // Add jump from Merge BB to the NormalDest. This is needed for the newly
-  // created direct invoke stmt -- as its NormalDst will be fixed up to MergeBB.
-  BranchInst::Create(II->getNormalDest(), *MergeBB);
-}
-
-// Find the PHI in BB that have the CallResult as the operand.
-static bool getCallRetPHINode(BasicBlock *BB, Instruction *Inst) {
-  BasicBlock *From = Inst->getParent();
-  for (auto &I : *BB) {
-    PHINode *PHI = dyn_cast<PHINode>(&I);
-    if (!PHI)
-      continue;
-    int IX = PHI->getBasicBlockIndex(From);
-    if (IX == -1)
-      continue;
-    Value *V = PHI->getIncomingValue(IX);
-    if (dyn_cast<Instruction>(V) == Inst)
-      return true;
-  }
-  return false;
-}
-
-// This method fixes up PHI nodes in BB where BB is the UnwindDest of an
-// invoke instruction. In BB, there may be PHIs with incoming block being
-// OrigBB (the MergeBB after if-then-else splitting). After moving the invoke
-// instructions to its own BB, OrigBB is no longer the predecessor block of BB.
-// Instead two new predecessors are added: IndirectCallBB and DirectCallBB,
-// so the PHI node's incoming BBs need to be fixed up accordingly.
-static void fixupPHINodeForUnwind(Instruction *Inst, BasicBlock *BB,
-                                  BasicBlock *OrigBB,
-                                  BasicBlock *IndirectCallBB,
-                                  BasicBlock *DirectCallBB) {
-  for (auto &I : *BB) {
-    PHINode *PHI = dyn_cast<PHINode>(&I);
-    if (!PHI)
-      continue;
-    int IX = PHI->getBasicBlockIndex(OrigBB);
-    if (IX == -1)
-      continue;
-    Value *V = PHI->getIncomingValue(IX);
-    PHI->addIncoming(V, IndirectCallBB);
-    PHI->setIncomingBlock(IX, DirectCallBB);
-  }
-}
-
-// This method fixes up PHI nodes in BB where BB is the NormalDest of an
-// invoke instruction. In BB, there may be PHIs with incoming block being
-// OrigBB (the MergeBB after if-then-else splitting). After moving the invoke
-// instructions to its own BB, a new incoming edge will be added to the original
-// NormalDstBB from the IndirectCallBB.
-static void fixupPHINodeForNormalDest(Instruction *Inst, BasicBlock *BB,
-                                      BasicBlock *OrigBB,
-                                      BasicBlock *IndirectCallBB,
-                                      Instruction *NewInst) {
-  for (auto &I : *BB) {
-    PHINode *PHI = dyn_cast<PHINode>(&I);
-    if (!PHI)
-      continue;
-    int IX = PHI->getBasicBlockIndex(OrigBB);
-    if (IX == -1)
-      continue;
-    Value *V = PHI->getIncomingValue(IX);
-    if (dyn_cast<Instruction>(V) == Inst) {
-      PHI->setIncomingBlock(IX, IndirectCallBB);
-      PHI->addIncoming(NewInst, OrigBB);
-      continue;
-    }
-    PHI->addIncoming(V, IndirectCallBB);
-  }
-}
-
-// Add a bitcast instruction to the direct-call return value if needed.
-static Instruction *insertCallRetCast(const Instruction *Inst,
-                                      Instruction *DirectCallInst,
-                                      Function *DirectCallee) {
-  if (Inst->getType()->isVoidTy())
-    return DirectCallInst;
-
-  Type *CallRetType = Inst->getType();
-  Type *FuncRetType = DirectCallee->getReturnType();
-  if (FuncRetType == CallRetType)
-    return DirectCallInst;
-
-  BasicBlock *InsertionBB;
-  if (CallInst *CI = dyn_cast<CallInst>(DirectCallInst))
-    InsertionBB = CI->getParent();
-  else
-    InsertionBB = (dyn_cast<InvokeInst>(DirectCallInst))->getNormalDest();
-
-  return (new BitCastInst(DirectCallInst, CallRetType, "",
-                          InsertionBB->getTerminator()));
-}
-
-// Create a DirectCall instruction in the DirectCallBB.
-// Parameter Inst is the indirect-call (invoke) instruction.
-// DirectCallee is the decl of the direct-call (invoke) target.
-// DirecallBB is the BB that the direct-call (invoke) instruction is inserted.
-// MergeBB is the bottom BB of the if-then-else-diamond after the
-// transformation. For invoke instruction, the edges from DirectCallBB and
-// IndirectCallBB to MergeBB are removed before this call (during
-// createIfThenElse). Stores the pointer to the Instruction that cast
-// the direct call in \p CastInst.
-static Instruction *createDirectCallInst(const Instruction *Inst,
-                                         Function *DirectCallee,
-                                         BasicBlock *DirectCallBB,
-                                         BasicBlock *MergeBB,
-                                         Instruction *&CastInst) {
-  Instruction *NewInst = Inst->clone();
-  if (CallInst *CI = dyn_cast<CallInst>(NewInst)) {
-    CI->setCalledFunction(DirectCallee);
-    CI->mutateFunctionType(DirectCallee->getFunctionType());
-  } else {
-    // Must be an invoke instruction. Direct invoke's normal destination is
-    // fixed up to MergeBB. MergeBB is the place where return cast is inserted.
-    // Also since IndirectCallBB does not have an edge to MergeBB, there is no
-    // need to insert new PHIs into MergeBB.
-    InvokeInst *II = dyn_cast<InvokeInst>(NewInst);
-    assert(II);
-    II->setCalledFunction(DirectCallee);
-    II->mutateFunctionType(DirectCallee->getFunctionType());
-    II->setNormalDest(MergeBB);
-  }
-
-  DirectCallBB->getInstList().insert(DirectCallBB->getFirstInsertionPt(),
-                                     NewInst);
-
-  // Clear the value profile data.
-  NewInst->setMetadata(LLVMContext::MD_prof, nullptr);
-  CallSite NewCS(NewInst);
-  FunctionType *DirectCalleeType = DirectCallee->getFunctionType();
-  unsigned ParamNum = DirectCalleeType->getFunctionNumParams();
-  for (unsigned I = 0; I < ParamNum; ++I) {
-    Type *ATy = NewCS.getArgument(I)->getType();
-    Type *PTy = DirectCalleeType->getParamType(I);
-    if (ATy != PTy) {
-      BitCastInst *BI = new BitCastInst(NewCS.getArgument(I), PTy, "", NewInst);
-      NewCS.setArgument(I, BI);
-    }
-  }
-
-  CastInst = insertCallRetCast(Inst, NewInst, DirectCallee);
-  return NewInst;
-}
-
-// Create a PHI to unify the return values of calls.
-static void insertCallRetPHI(Instruction *Inst, Instruction *CallResult,
-                             Function *DirectCallee) {
-  if (Inst->getType()->isVoidTy())
-    return;
-
-  if (Inst->use_empty())
-    return;
-
-  BasicBlock *RetValBB = CallResult->getParent();
-
-  BasicBlock *PHIBB;
-  if (InvokeInst *II = dyn_cast<InvokeInst>(CallResult))
-    RetValBB = II->getNormalDest();
-
-  PHIBB = RetValBB->getSingleSuccessor();
-  if (getCallRetPHINode(PHIBB, Inst))
-    return;
-
-  PHINode *CallRetPHI = PHINode::Create(Inst->getType(), 0);
-  PHIBB->getInstList().push_front(CallRetPHI);
-  Inst->replaceAllUsesWith(CallRetPHI);
-  CallRetPHI->addIncoming(Inst, Inst->getParent());
-  CallRetPHI->addIncoming(CallResult, RetValBB);
-}
-
-// This function does the actual indirect-call promotion transformation:
-// For an indirect-call like:
-//     Ret = (*Foo)(Args);
-// It transforms to:
-//     if (Foo == DirectCallee)
-//        Ret1 = DirectCallee(Args);
-//     else
-//        Ret2 = (*Foo)(Args);
-//     Ret = phi(Ret1, Ret2);
-// It adds type casts for the args do not match the parameters and the return
-// value. Branch weights metadata also updated.
-// If \p AttachProfToDirectCall is true, a prof metadata is attached to the
-// new direct call to contain \p Count. This is used by SamplePGO inliner to
-// check callsite hotness.
-// Returns the promoted direct call instruction.
-Instruction *llvm::promoteIndirectCall(Instruction *Inst,
-                                       Function *DirectCallee, uint64_t Count,
-                                       uint64_t TotalCount,
-                                       bool AttachProfToDirectCall,
-                                       OptimizationRemarkEmitter *ORE) {
-  assert(DirectCallee != nullptr);
-  BasicBlock *BB = Inst->getParent();
-  // Just to suppress the non-debug build warning.
-  (void)BB;
-  DEBUG(dbgs() << "\n\n== Basic Block Before ==\n");
-  DEBUG(dbgs() << *BB << "\n");
-
-  BasicBlock *DirectCallBB, *IndirectCallBB, *MergeBB;
-  createIfThenElse(Inst, DirectCallee, Count, TotalCount, &DirectCallBB,
-                   &IndirectCallBB, &MergeBB);
-
-  // If the return type of the NewInst is not the same as the Inst, a CastInst
-  // is needed for type casting. Otherwise CastInst is the same as NewInst.
-  Instruction *CastInst = nullptr;
   Instruction *NewInst =
-      createDirectCallInst(Inst, DirectCallee, DirectCallBB, MergeBB, CastInst);
+      promoteCallWithIfThenElse(CallSite(Inst), DirectCallee, BranchWeights);
 
   if (AttachProfToDirectCall) {
     SmallVector<uint32_t, 1> Weights;
@@ -592,33 +316,6 @@
     NewInst->setMetadata(LLVMContext::MD_prof, MDB.createBranchWeights(Weights));
   }
 
-  // Move Inst from MergeBB to IndirectCallBB.
-  Inst->removeFromParent();
-  IndirectCallBB->getInstList().insert(IndirectCallBB->getFirstInsertionPt(),
-                                       Inst);
-
-  if (InvokeInst *II = dyn_cast<InvokeInst>(Inst)) {
-    // At this point, the original indirect invoke instruction has the original
-    // UnwindDest and NormalDest. For the direct invoke instruction, the
-    // NormalDest points to MergeBB, and MergeBB jumps to the original
-    // NormalDest. MergeBB might have a new bitcast instruction for the return
-    // value. The PHIs are with the original NormalDest. Since we now have two
-    // incoming edges to NormalDest and UnwindDest, we have to do some fixups.
-    //
-    // UnwindDest will not use the return value. So pass nullptr here.
-    fixupPHINodeForUnwind(Inst, II->getUnwindDest(), MergeBB, IndirectCallBB,
-                          DirectCallBB);
-    // We don't need to update the operand from NormalDest for DirectCallBB.
-    // Pass nullptr here.
-    fixupPHINodeForNormalDest(Inst, II->getNormalDest(), MergeBB,
-                              IndirectCallBB, CastInst);
-  }
-
-  insertCallRetPHI(Inst, CastInst, DirectCallee);
-
-  DEBUG(dbgs() << "\n== Basic Blocks After ==\n");
-  DEBUG(dbgs() << *BB << *DirectCallBB << *IndirectCallBB << *MergeBB << "\n");
-
   using namespace ore;
 
   if (ORE)
@@ -639,8 +336,8 @@
 
   for (auto &C : Candidates) {
     uint64_t Count = C.Count;
-    promoteIndirectCall(Inst, C.TargetFunction, Count, TotalCount, SamplePGO,
-                        &ORE);
+    pgo::promoteIndirectCall(Inst, C.TargetFunction, Count, TotalCount,
+                             SamplePGO, &ORE);
     assert(TotalCount >= Count);
     TotalCount -= Count;
     NumOfPGOICallPromotion++;