[WPD] Fix incorrect devirtualization after indirect call promotion

[tejohnson]

Summary

Add a dominance check to ensure that the possible devirtualizable
call is actually dominated by the type test/checked load intrinsic being
analyzed. With PGO, after indirect call promotion is performed during
the compile step, followed by inlining, we may have a type test in the
promoted and inlined sequence that allows an indirect call in that
sequence to be devirtualized. That indirect call (inserted by inlining
after promotion) will share the same vtable pointer as the fallback
indirect call that cannot be devirtualized.

Before this patch the code was incorrectly devirtualizing the fallback
indirect call.

See the new test and the example described there for more details.

include/llvm/Analysis/TypeMetadataUtils.h

 #ifndef LLVM_ANALYSIS_TYPEMETADATAUTILS_H
 #define LLVM_ANALYSIS_TYPEMETADATAUTILS_H
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/CallSite.h"
 
 namespace llvm {
 
+class DominatorTree;
+
 /// The type of CFI jumptable needed for a function.
 enum CfiFunctionLinkage {
   CFL_Definition = 0,
   CFL_Declaration = 1,
   CFL_WeakDeclaration = 2
 };
 
 /// A call site that could be devirtualized.
 struct DevirtCallSite {
   /// The offset from the address point to the virtual function.
   uint64_t Offset;
   /// The call site itself.
   CallSite CS;
 };
 
 /// Given a call to the intrinsic \@llvm.type.test, find all devirtualizable
 /// call sites based on the call and return them in DevirtCalls.
 void findDevirtualizableCallsForTypeTest(
     SmallVectorImpl<DevirtCallSite> &DevirtCalls,
-    SmallVectorImpl<CallInst *> &Assumes, const CallInst *CI);
+    SmallVectorImpl<CallInst *> &Assumes, const CallInst *CI,
+    DominatorTree &DT);
 
 /// Given a call to the intrinsic \@llvm.type.checked.load, find all
 /// devirtualizable call sites based on the call and return them in DevirtCalls.
 void findDevirtualizableCallsForTypeCheckedLoad(
     SmallVectorImpl<DevirtCallSite> &DevirtCalls,
     SmallVectorImpl<Instruction *> &LoadedPtrs,
     SmallVectorImpl<Instruction *> &Preds, bool &HasNonCallUses,
-    const CallInst *CI);
+    const CallInst *CI, DominatorTree &DT);
 }
 
 #endif

lib/Analysis/ModuleSummaryAnalysis.cpp

 
 /// If this intrinsic call requires that we add information to the function
 /// summary, do so via the non-constant reference arguments.
 static void addIntrinsicToSummary(
     const CallInst *CI, SetVector<GlobalValue::GUID> &TypeTests,
     SetVector<FunctionSummary::VFuncId> &TypeTestAssumeVCalls,
     SetVector<FunctionSummary::VFuncId> &TypeCheckedLoadVCalls,
     SetVector<FunctionSummary::ConstVCall> &TypeTestAssumeConstVCalls,
-    SetVector<FunctionSummary::ConstVCall> &TypeCheckedLoadConstVCalls) {
+    SetVector<FunctionSummary::ConstVCall> &TypeCheckedLoadConstVCalls,
+    DominatorTree &DT) {
   switch (CI->getCalledFunction()->getIntrinsicID()) {
   case Intrinsic::type_test: {
     auto *TypeMDVal = cast<MetadataAsValue>(CI->getArgOperand(1));
     auto *TypeId = dyn_cast<MDString>(TypeMDVal->getMetadata());
     if (!TypeId)
       break;
     GlobalValue::GUID Guid = GlobalValue::getGUID(TypeId->getString());
 
     // Produce a summary from type.test intrinsics. We only summarize type.test
     // intrinsics that are used other than by an llvm.assume intrinsic.
     // Intrinsics that are assumed are relevant only to the devirtualization
     // pass, not the type test lowering pass.
     bool HasNonAssumeUses = llvm::any_of(CI->uses(), [](const Use &CIU) {
       auto *AssumeCI = dyn_cast<CallInst>(CIU.getUser());
       if (!AssumeCI)
         return true;
       Function *F = AssumeCI->getCalledFunction();
       return !F || F->getIntrinsicID() != Intrinsic::assume;
     });
     if (HasNonAssumeUses)
       TypeTests.insert(Guid);
 
     SmallVector<DevirtCallSite, 4> DevirtCalls;
     SmallVector<CallInst *, 4> Assumes;
-    findDevirtualizableCallsForTypeTest(DevirtCalls, Assumes, CI);
+    findDevirtualizableCallsForTypeTest(DevirtCalls, Assumes, CI, DT);
     for (auto &Call : DevirtCalls)
       addVCallToSet(Call, Guid, TypeTestAssumeVCalls,
                     TypeTestAssumeConstVCalls);
 
     break;
   }
 
   case Intrinsic::type_checked_load: {
     auto *TypeMDVal = cast<MetadataAsValue>(CI->getArgOperand(2));
     auto *TypeId = dyn_cast<MDString>(TypeMDVal->getMetadata());
     if (!TypeId)
       break;
     GlobalValue::GUID Guid = GlobalValue::getGUID(TypeId->getString());
 
     SmallVector<DevirtCallSite, 4> DevirtCalls;
     SmallVector<Instruction *, 4> LoadedPtrs;
     SmallVector<Instruction *, 4> Preds;
     bool HasNonCallUses = false;
     findDevirtualizableCallsForTypeCheckedLoad(DevirtCalls, LoadedPtrs, Preds,
-                                               HasNonCallUses, CI);
+                                               HasNonCallUses, CI, DT);
     // Any non-call uses of the result of llvm.type.checked.load will
     // prevent us from optimizing away the llvm.type.test.
     if (HasNonCallUses)
       TypeTests.insert(Guid);
     for (auto &Call : DevirtCalls)
       addVCallToSet(Call, Guid, TypeCheckedLoadVCalls,
                     TypeCheckedLoadConstVCalls);
 
     break;
   }
   default:
     break;
   }
 }
 
-static void
-computeFunctionSummary(ModuleSummaryIndex &Index, const Module &M,
-                       const Function &F, BlockFrequencyInfo *BFI,
-                       ProfileSummaryInfo *PSI, bool HasLocalsInUsedOrAsm,
-                       DenseSet<GlobalValue::GUID> &CantBePromoted) {
+static void computeFunctionSummary(
+    ModuleSummaryIndex &Index, const Module &M, const Function &F,
+    BlockFrequencyInfo *BFI, ProfileSummaryInfo *PSI, DominatorTree &DT,
+    bool HasLocalsInUsedOrAsm, DenseSet<GlobalValue::GUID> &CantBePromoted) {
   // Summary not currently supported for anonymous functions, they should
   // have been named.
   assert(F.hasName());
 
   unsigned NumInsts = 0;
   // Map from callee ValueId to profile count. Used to accumulate profile
   // counts for all static calls to a given callee.
   MapVector<ValueInfo, CalleeInfo> CallGraphEdges;
for (const Instruction &I : BB) {
         CalledFunction = dyn_cast<Function>(GA->getBaseObject());
       }
       // Check if this is a direct call to a known function or a known
       // intrinsic, or an indirect call with profile data.
       if (CalledFunction) {
         if (CI && CalledFunction->isIntrinsic()) {
           addIntrinsicToSummary(
               CI, TypeTests, TypeTestAssumeVCalls, TypeCheckedLoadVCalls,
-              TypeTestAssumeConstVCalls, TypeCheckedLoadConstVCalls);
+              TypeTestAssumeConstVCalls, TypeCheckedLoadConstVCalls, DT);
           continue;
         }
         // We should have named any anonymous globals
         assert(CalledFunction->hasName());
         auto ScaledCount = PSI->getProfileCount(&I, BFI);
         auto Hotness = ScaledCount ? getHotness(ScaledCount.getValue(), PSI)
                                    : CalleeInfo::HotnessType::Unknown;
         if (ForceSummaryEdgesCold != FunctionSummary::FSHT_None)
ModuleSummaryIndex llvm::buildModuleSummaryIndex(
   }
 
   // Compute summaries for all functions defined in module, and save in the
   // index.
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
 
+    DominatorTree DT(const_cast<Function &>(F));
     BlockFrequencyInfo *BFI = nullptr;
     std::unique_ptr<BlockFrequencyInfo> BFIPtr;
     if (GetBFICallback)
       BFI = GetBFICallback(F);
     else if (F.hasProfileData()) {
-      LoopInfo LI{DominatorTree(const_cast<Function &>(F))};
+      LoopInfo LI{DT};
       BranchProbabilityInfo BPI{F, LI};
       BFIPtr = llvm::make_unique<BlockFrequencyInfo>(F, BPI, LI);
       BFI = BFIPtr.get();
     }
 
-    computeFunctionSummary(Index, M, F, BFI, PSI,
+    computeFunctionSummary(Index, M, F, BFI, PSI, DT,
                            !LocalsUsed.empty() || HasLocalInlineAsmSymbol,
                            CantBePromoted);
   }
 
   // Compute summaries for all variables defined in module, and save in the
   // index.
   for (const GlobalVariable &G : M.globals()) {
     if (G.isDeclaration())

lib/Analysis/TypeMetadataUtils.cpp

 //===- TypeMetadataUtils.cpp - Utilities related to type metadata ---------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains functions that make it easier to manipulate type metadata
 // for devirtualization.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/TypeMetadataUtils.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Module.h"
 
 using namespace llvm;
 
 // Search for virtual calls that call FPtr and add them to DevirtCalls.
 static void
 findCallsAtConstantOffset(SmallVectorImpl<DevirtCallSite> &DevirtCalls,
-                          bool *HasNonCallUses, Value *FPtr, uint64_t Offset) {
+                          bool *HasNonCallUses, Value *FPtr, uint64_t Offset,
+                          const CallInst *CI, DominatorTree &DT) {
   for (const Use &U : FPtr->uses()) {
-    Value *User = U.getUser();
+    Instruction *User = cast<Instruction>(U.getUser());
+    // Ignore this instruction if it is not dominated by the type intrinsic
+    // being analyzed. Otherwise we may transform a call sharing the same
+    // vtable pointer incorrectly. Specifically, this situation can arise
+    // after indirect call promotion and inlining, where we may have uses
+    // of the vtable pointer guarded by a function pointer check, and a fallback
+    // indirect call.
+    if (!DT.dominates(CI, User))
+      continue;
     if (isa<BitCastInst>(User)) {
-      findCallsAtConstantOffset(DevirtCalls, HasNonCallUses, User, Offset);
+      findCallsAtConstantOffset(DevirtCalls, HasNonCallUses, User, Offset, CI,
+                                DT);
     } else if (auto CI = dyn_cast<CallInst>(User)) {
       DevirtCalls.push_back({Offset, CI});
     } else if (auto II = dyn_cast<InvokeInst>(User)) {
       DevirtCalls.push_back({Offset, II});
     } else if (HasNonCallUses) {
       *HasNonCallUses = true;
     }
   }
 }
 
 // Search for virtual calls that load from VPtr and add them to DevirtCalls.
-static void
-findLoadCallsAtConstantOffset(const Module *M,
-                              SmallVectorImpl<DevirtCallSite> &DevirtCalls,
-                              Value *VPtr, int64_t Offset) {
+static void findLoadCallsAtConstantOffset(
+    const Module *M, SmallVectorImpl<DevirtCallSite> &DevirtCalls, Value *VPtr,
+    int64_t Offset, const CallInst *CI, DominatorTree &DT) {
   for (const Use &U : VPtr->uses()) {
     Value *User = U.getUser();
     if (isa<BitCastInst>(User)) {
-      findLoadCallsAtConstantOffset(M, DevirtCalls, User, Offset);
+      findLoadCallsAtConstantOffset(M, DevirtCalls, User, Offset, CI, DT);
     } else if (isa<LoadInst>(User)) {
-      findCallsAtConstantOffset(DevirtCalls, nullptr, User, Offset);
+      findCallsAtConstantOffset(DevirtCalls, nullptr, User, Offset, CI, DT);
     } else if (auto GEP = dyn_cast<GetElementPtrInst>(User)) {
       // Take into account the GEP offset.
       if (VPtr == GEP->getPointerOperand() && GEP->hasAllConstantIndices()) {
         SmallVector<Value *, 8> Indices(GEP->op_begin() + 1, GEP->op_end());
         int64_t GEPOffset = M->getDataLayout().getIndexedOffsetInType(
             GEP->getSourceElementType(), Indices);
-        findLoadCallsAtConstantOffset(M, DevirtCalls, User, Offset + GEPOffset);
+        findLoadCallsAtConstantOffset(M, DevirtCalls, User, Offset + GEPOffset,
+                                      CI, DT);
       }
     }
   }
 }
 
 void llvm::findDevirtualizableCallsForTypeTest(
     SmallVectorImpl<DevirtCallSite> &DevirtCalls,
-    SmallVectorImpl<CallInst *> &Assumes, const CallInst *CI) {
+    SmallVectorImpl<CallInst *> &Assumes, const CallInst *CI,
+    DominatorTree &DT) {
   assert(CI->getCalledFunction()->getIntrinsicID() == Intrinsic::type_test);
 
   const Module *M = CI->getParent()->getParent()->getParent();
 
   // Find llvm.assume intrinsics for this llvm.type.test call.
   for (const Use &CIU : CI->uses()) {
     if (auto *AssumeCI = dyn_cast<CallInst>(CIU.getUser())) {
       Function *F = AssumeCI->getCalledFunction();
       if (F && F->getIntrinsicID() == Intrinsic::assume)
         Assumes.push_back(AssumeCI);
     }
   }
 
   // If we found any, search for virtual calls based on %p and add them to
   // DevirtCalls.
   if (!Assumes.empty())
-    findLoadCallsAtConstantOffset(M, DevirtCalls,
-                                  CI->getArgOperand(0)->stripPointerCasts(), 0);
+    findLoadCallsAtConstantOffset(
+        M, DevirtCalls, CI->getArgOperand(0)->stripPointerCasts(), 0, CI, DT);
 }
 
 void llvm::findDevirtualizableCallsForTypeCheckedLoad(
     SmallVectorImpl<DevirtCallSite> &DevirtCalls,
     SmallVectorImpl<Instruction *> &LoadedPtrs,
     SmallVectorImpl<Instruction *> &Preds, bool &HasNonCallUses,
-    const CallInst *CI) {
+    const CallInst *CI, DominatorTree &DT) {
   assert(CI->getCalledFunction()->getIntrinsicID() ==
          Intrinsic::type_checked_load);
 
   auto *Offset = dyn_cast<ConstantInt>(CI->getArgOperand(1));
   if (!Offset) {
     HasNonCallUses = true;
     return;
   }
if (auto EVI = dyn_cast<ExtractValueInst>(CIU)) {
         continue;
       }
     }
     HasNonCallUses = true;
   }
 
   for (Value *LoadedPtr : LoadedPtrs)
     findCallsAtConstantOffset(DevirtCalls, &HasNonCallUses, LoadedPtr,
-                              Offset->getZExtValue());
+                              Offset->getZExtValue(), CI, DT);
 }

lib/Transforms/IPO/WholeProgramDevirt.cpp

 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/TypeMetadataUtils.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
void VTableSlotInfo::addCallSite(Value *VTable, CallSite CS,
   auto &CSI = findCallSiteInfo(CS);
   CSI.AllCallSitesDevirted = false;
   CSI.CallSites.push_back({VTable, CS, NumUnsafeUses});
 }
 
 struct DevirtModule {
   Module &M;
   function_ref<AAResults &(Function &)> AARGetter;
+  function_ref<DominatorTree &(Function &)> LookupDomTree;
 
   ModuleSummaryIndex *ExportSummary;
   const ModuleSummaryIndex *ImportSummary;
 
   IntegerType *Int8Ty;
   PointerType *Int8PtrTy;
   IntegerType *Int32Ty;
   IntegerType *Int64Ty;
struct DevirtModule {
   // devirtualizing it or by applying virtual constant propagation), we
   // decrement the value stored in this map. If a value reaches zero, we can
   // eliminate the type check by RAUWing the associated llvm.type.test call with
   // true.
   std::map<CallInst *, unsigned> NumUnsafeUsesForTypeTest;
 
   DevirtModule(Module &M, function_ref<AAResults &(Function &)> AARGetter,
                function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter,
+               function_ref<DominatorTree &(Function &)> LookupDomTree,
                ModuleSummaryIndex *ExportSummary,
                const ModuleSummaryIndex *ImportSummary)
-      : M(M), AARGetter(AARGetter), ExportSummary(ExportSummary),
-        ImportSummary(ImportSummary), Int8Ty(Type::getInt8Ty(M.getContext())),
+      : M(M), AARGetter(AARGetter), LookupDomTree(LookupDomTree),
+        ExportSummary(ExportSummary), ImportSummary(ImportSummary),
+        Int8Ty(Type::getInt8Ty(M.getContext())),
         Int8PtrTy(Type::getInt8PtrTy(M.getContext())),
         Int32Ty(Type::getInt32Ty(M.getContext())),
         Int64Ty(Type::getInt64Ty(M.getContext())),
         IntPtrTy(M.getDataLayout().getIntPtrType(M.getContext(), 0)),
         RemarksEnabled(areRemarksEnabled()), OREGetter(OREGetter) {
     assert(!(ExportSummary && ImportSummary));
   }
 
struct DevirtModule {
   // If we were able to eliminate all unsafe uses for a type checked load,
   // eliminate the associated type tests by replacing them with true.
   void removeRedundantTypeTests();
 
   bool run();
 
   // Lower the module using the action and summary passed as command line
   // arguments. For testing purposes only.
-  static bool runForTesting(
-      Module &M, function_ref<AAResults &(Function &)> AARGetter,
-      function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter);
+  static bool
+  runForTesting(Module &M, function_ref<AAResults &(Function &)> AARGetter,
+                function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter,
+                function_ref<DominatorTree &(Function &)> LookupDomTree);
 };
 
 struct WholeProgramDevirt : public ModulePass {
   static char ID;
 
   bool UseCommandLine = false;
 
   ModuleSummaryIndex *ExportSummary;
bool runOnModule(Module &M) override {
     // In the old pass manager, this is harder, so we just build
     // an optimization remark emitter on the fly, when we need it.
     std::unique_ptr<OptimizationRemarkEmitter> ORE;
     auto OREGetter = [&](Function *F) -> OptimizationRemarkEmitter & {
       ORE = make_unique<OptimizationRemarkEmitter>(F);
       return *ORE;
     };
 
+    auto LookupDomTree = [this](Function &F) -> DominatorTree & {
+      return this->getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
+    };
+
     if (UseCommandLine)
-      return DevirtModule::runForTesting(M, LegacyAARGetter(*this), OREGetter);
+      return DevirtModule::runForTesting(M, LegacyAARGetter(*this), OREGetter,
+                                         LookupDomTree);
 
-    return DevirtModule(M, LegacyAARGetter(*this), OREGetter, ExportSummary,
-                        ImportSummary)
+    return DevirtModule(M, LegacyAARGetter(*this), OREGetter, LookupDomTree,
+                        ExportSummary, ImportSummary)
         .run();
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
+    AU.addRequired<DominatorTreeWrapperPass>();
   }
 };
 
 } // end anonymous namespace
 
 INITIALIZE_PASS_BEGIN(WholeProgramDevirt, "wholeprogramdevirt",
                       "Whole program devirtualization", false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_END(WholeProgramDevirt, "wholeprogramdevirt",
                     "Whole program devirtualization", false, false)
 char WholeProgramDevirt::ID = 0;
 
 ModulePass *
 llvm::createWholeProgramDevirtPass(ModuleSummaryIndex *ExportSummary,
                                    const ModuleSummaryIndex *ImportSummary) {
   return new WholeProgramDevirt(ExportSummary, ImportSummary);
 }
 
 PreservedAnalyses WholeProgramDevirtPass::run(Module &M,
                                               ModuleAnalysisManager &AM) {
   auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   auto AARGetter = [&](Function &F) -> AAResults & {
     return FAM.getResult<AAManager>(F);
   };
   auto OREGetter = [&](Function *F) -> OptimizationRemarkEmitter & {
     return FAM.getResult<OptimizationRemarkEmitterAnalysis>(*F);
   };
-  if (!DevirtModule(M, AARGetter, OREGetter, ExportSummary, ImportSummary)
+  auto LookupDomTree = [&FAM](Function &F) -> DominatorTree & {
+    return FAM.getResult<DominatorTreeAnalysis>(F);
+  };
+  if (!DevirtModule(M, AARGetter, OREGetter, LookupDomTree, ExportSummary,
+                    ImportSummary)
            .run())
     return PreservedAnalyses::all();
   return PreservedAnalyses::none();
 }
 
 bool DevirtModule::runForTesting(
     Module &M, function_ref<AAResults &(Function &)> AARGetter,
-    function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter) {
+    function_ref<OptimizationRemarkEmitter &(Function *)> OREGetter,
+    function_ref<DominatorTree &(Function &)> LookupDomTree) {
   ModuleSummaryIndex Summary(/*HaveGVs=*/false);
 
   // Handle the command-line summary arguments. This code is for testing
   // purposes only, so we handle errors directly.
   if (!ClReadSummary.empty()) {
     ExitOnError ExitOnErr("-wholeprogramdevirt-read-summary: " + ClReadSummary +
                           ": ");
     auto ReadSummaryFile =
         ExitOnErr(errorOrToExpected(MemoryBuffer::getFile(ClReadSummary)));
 
     yaml::Input In(ReadSummaryFile->getBuffer());
     In >> Summary;
     ExitOnErr(errorCodeToError(In.error()));
   }
 
   bool Changed =
       DevirtModule(
-          M, AARGetter, OREGetter,
+          M, AARGetter, OREGetter, LookupDomTree,
           ClSummaryAction == PassSummaryAction::Export ? &Summary : nullptr,
           ClSummaryAction == PassSummaryAction::Import ? &Summary : nullptr)
           .run();
 
   if (!ClWriteSummary.empty()) {
     ExitOnError ExitOnErr(
         "-wholeprogramdevirt-write-summary: " + ClWriteSummary + ": ");
     std::error_code EC;

 
 void DevirtModule::scanTypeTestUsers(Function *TypeTestFunc,
                                      Function *AssumeFunc) {
   // Find all virtual calls via a virtual table pointer %p under an assumption
   // of the form llvm.assume(llvm.type.test(%p, %md)). This indicates that %p
   // points to a member of the type identifier %md. Group calls by (type ID,
   // offset) pair (effectively the identity of the virtual function) and store
   // to CallSlots.
-  DenseSet<Value *> SeenPtrs;
+  DenseSet<CallSite> SeenCallSites;
   for (auto I = TypeTestFunc->use_begin(), E = TypeTestFunc->use_end();
        I != E;) {
     auto CI = dyn_cast<CallInst>(I->getUser());
     ++I;
     if (!CI)
       continue;
 
     // Search for virtual calls based on %p and add them to DevirtCalls.
     SmallVector<DevirtCallSite, 1> DevirtCalls;
     SmallVector<CallInst *, 1> Assumes;
-    findDevirtualizableCallsForTypeTest(DevirtCalls, Assumes, CI);
+    auto &DT = LookupDomTree(*CI->getFunction());
+    findDevirtualizableCallsForTypeTest(DevirtCalls, Assumes, CI, DT);
 
-    // If we found any, add them to CallSlots. Only do this if we haven't seen
-    // the vtable pointer before, as it may have been CSE'd with pointers from
-    // other call sites, and we don't want to process call sites multiple times.
+    // If we found any, add them to CallSlots.
     if (!Assumes.empty()) {
       Metadata *TypeId =
           cast<MetadataAsValue>(CI->getArgOperand(1))->getMetadata();
       Value *Ptr = CI->getArgOperand(0)->stripPointerCasts();
-      if (SeenPtrs.insert(Ptr).second) {
-        for (DevirtCallSite Call : DevirtCalls) {
+        // Only add this CallSite if we haven't seen it before. The vtable
+        // pointer may have been CSE'd with pointers from other call sites,
+        // and we don't want to process call sites multiple times. We can't
+        // just skip the vtable Ptr if it has been seen before, however, since
+        // it may be shared by type tests that dominate different calls.
+        if (SeenCallSites.insert(Call.CS).second)
           CallSlots[{TypeId, Call.Offset}].addCallSite(Ptr, Call.CS, nullptr);
-        }
-      }
-    }
 
     // We no longer need the assumes or the type test.
     for (auto Assume : Assumes)
       Assume->eraseFromParent();
     // We can't use RecursivelyDeleteTriviallyDeadInstructions here because we
     // may use the vtable argument later.
     if (CI->use_empty())
       CI->eraseFromParent();
for (auto I = TypeCheckedLoadFunc->use_begin(),
     Value *Offset = CI->getArgOperand(1);
     Value *TypeIdValue = CI->getArgOperand(2);
     Metadata *TypeId = cast<MetadataAsValue>(TypeIdValue)->getMetadata();
 
     SmallVector<DevirtCallSite, 1> DevirtCalls;
     SmallVector<Instruction *, 1> LoadedPtrs;
     SmallVector<Instruction *, 1> Preds;
     bool HasNonCallUses = false;
+    auto &DT = LookupDomTree(*CI->getFunction());
     findDevirtualizableCallsForTypeCheckedLoad(DevirtCalls, LoadedPtrs, Preds,
-                                               HasNonCallUses, CI);
+                                               HasNonCallUses, CI, DT);
 
     // Start by generating "pessimistic" code that explicitly loads the function
     // pointer from the vtable and performs the type check. If possible, we will
     // eliminate the load and the type check later.
 
     // If possible, only generate the load at the point where it is used.
     // This helps avoid unnecessary spills.
     IRBuilder<> LoadB(

test/ThinLTO/X86/devirt-after-icp.ll

+; REQUIRES: x86-registered-target
+
+; Test devirtualization through the thin link and backend, ensuring that
+; it is only applied when the type test corresponding to a devirtualization
+; dominates an indirect call using the same vtable pointer. Indirect
+; call promotion and inlining may introduce a guarded indirect call
+; that can be promoted, which uses the same vtable address as the fallback
+; indirect call that cannot be devirtualized.
+
+; The code below illustrates the structure when we started with code like:
+;
+; class A {
+;  public:
+;   virtual int foo() { return 1; }
+;   virtual int bar() { return 1; }
+; };
+; class B : public A {
+;  public:
+;   virtual int foo();
+;   virtual int bar();
+; };
+;
+; int foo(A *a) {
+;   return a->foo(); // ICP profile says most calls are to B::foo()
+; }
+;
+; int B::foo() {
+;   return bar();
+; }
+;
+; After the compile step, which will perform ICP and a round of inlining, we
+; have something like:
+; int foo(A *a) {
+;   if (&a->foo() == B::foo())
+;     return ((B*)a)->bar(); // Inlined from promoted direct call to B::foo()
+;   else
+;     return a->foo();
+;
+; The inlined code seqence will have a type test against "_ZTS1B",
+; which will allow us to devirtualize indirect call ((B*)a)->bar() to B::bar();
+; Both that type test and the one for the fallback a->foo() indirect call
+; will use the same vtable pointer. Without a dominance check, we could
+; incorrectly devirtualize a->foo() to B::foo();
+
+; RUN: opt -thinlto-bc -o %t.o %s
+
+; Legacy PM
+; RUN: llvm-lto2 run %t.o -save-temps -pass-remarks=. \
+; RUN:   -o %t3 \
+; RUN:   -r=%t.o,_Z3bazP1A,px \
+; RUN:   -r=%t.o,_ZN1A3fooEv, \
+; RUN:   -r=%t.o,_ZN1A3barEv, \
+; RUN:   -r=%t.o,_ZN1B3fooEv, \
+; RUN:   -r=%t.o,_ZN1B3barEv, \
+; RUN:   -r=%t.o,_ZTV1A, \
+; RUN:   -r=%t.o,_ZTV1B, \
+; RUN:   -r=%t.o,_ZN1A3fooEv, \
+; RUN:   -r=%t.o,_ZN1A3barEv, \
+; RUN:   -r=%t.o,_ZN1B3fooEv, \
+; RUN:   -r=%t.o,_ZN1B3barEv, \
+; RUN:   -r=%t.o,_ZTV1A,px \
+; RUN:   -r=%t.o,_ZTV1B,px 2>&1 | FileCheck %s --check-prefix=REMARK
+; RUN: llvm-dis %t3.1.4.opt.bc -o - | FileCheck %s --check-prefix=CHECK-IR
+
+; New PM
+; RUN: llvm-lto2 run %t.o -save-temps -use-new-pm -pass-remarks=. \
+; RUN:   -o %t3 \
+; RUN:   -r=%t.o,_Z3bazP1A,px \
+; RUN:   -r=%t.o,_ZN1A3fooEv, \
+; RUN:   -r=%t.o,_ZN1A3barEv, \
+; RUN:   -r=%t.o,_ZN1B3fooEv, \
+; RUN:   -r=%t.o,_ZN1B3barEv, \
+; RUN:   -r=%t.o,_ZTV1A, \
+; RUN:   -r=%t.o,_ZTV1B, \
+; RUN:   -r=%t.o,_ZN1A3fooEv, \
+; RUN:   -r=%t.o,_ZN1A3barEv, \
+; RUN:   -r=%t.o,_ZN1B3fooEv, \
+; RUN:   -r=%t.o,_ZN1B3barEv, \
+; RUN:   -r=%t.o,_ZTV1A,px \
+; RUN:   -r=%t.o,_ZTV1B,px 2>&1 | FileCheck %s --check-prefix=REMARK
+; RUN: llvm-dis %t3.1.4.opt.bc -o - | FileCheck %s --check-prefix=CHECK-IR
+
+; We should only devirtualize the inlined call to bar().
+; REMARK-NOT: single-impl: devirtualized a call to _ZN1B3fooEv
+; REMARK: single-impl: devirtualized a call to _ZN1B3barEv
+; REMARK-NOT: single-impl: devirtualized a call to _ZN1B3fooEv
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+%class.A = type { i32 (...)** }
+%class.B = type { %class.A }
+
+@_ZTV1A = linkonce_odr hidden unnamed_addr constant { [4 x i8*] } { [4 x i8*] [i8* null, i8* undef, i8* bitcast (i32 (%class.A*)* @_ZN1A3fooEv to i8*), i8* bitcast (i32 (%class.A*)* @_ZN1A3barEv to i8*)] }, align 8, !type !0
+@_ZTV1B = hidden unnamed_addr constant { [4 x i8*] } { [4 x i8*] [i8* null, i8* undef, i8* bitcast (i32 (%class.B*)* @_ZN1B3fooEv to i8*), i8* bitcast (i32 (%class.B*)* @_ZN1B3barEv to i8*)] }, align 8, !type !0, !type !1
+
+define hidden i32 @_Z3bazP1A(%class.A* %a) local_unnamed_addr {
+entry:
+  %0 = bitcast %class.A* %a to i32 (%class.A*)***
+  %vtable = load i32 (%class.A*)**, i32 (%class.A*)*** %0, align 8
+  %1 = bitcast i32 (%class.A*)** %vtable to i8*
+  %2 = tail call i1 @llvm.type.test(i8* %1, metadata !"_ZTS1A")
+  tail call void @llvm.assume(i1 %2)
+  %3 = load i32 (%class.A*)*, i32 (%class.A*)** %vtable, align 8
+  ; This is the compare instruction inserted by ICP
+  %4 = icmp eq i32 (%class.A*)* %3, bitcast (i32 (%class.B*)* @_ZN1B3fooEv to i32 (%class.A*)*)
+  br i1 %4, label %if.true.direct_targ, label %if.false.orig_indirect
+
+; This block contains the promoted and inlined call to B::foo();
+; CHECK-IR: if.true.direct_targ:                              ; preds = %entry
+if.true.direct_targ:                              ; preds = %entry
+  %5 = bitcast %class.A* %a to %class.B*
+  %6 = bitcast i32 (%class.A*)** %vtable to i8*
+  %7 = tail call i1 @llvm.type.test(i8* %6, metadata !"_ZTS1B")
+  tail call void @llvm.assume(i1 %7)
+  %vfn.i1 = getelementptr inbounds i32 (%class.A*)*, i32 (%class.A*)** %vtable, i64 1
+  %vfn.i = bitcast i32 (%class.A*)** %vfn.i1 to i32 (%class.B*)**
+  %8 = load i32 (%class.B*)*, i32 (%class.B*)** %vfn.i, align 8
+; Call to bar() can be devirtualized to call to B::bar(), since it was
+; inlined from B::foo() after ICP introduced the guarded promotion.
+; CHECK-IR: %call.i = tail call i32 @_ZN1B3barEv(%class.B* %3)
+  %call.i = tail call i32 %8(%class.B* %5)
+  br label %if.end.icp
+
+; This block contains the fallback indirect call a->foo()
+; CHECK-IR: if.false.orig_indirect:
+if.false.orig_indirect:                           ; preds = %entry
+; Fallback indirect call to foo() cannot be devirtualized.
+; CHECK-IR: %call = tail call i32 %
+  %call = tail call i32 %3(%class.A* nonnull %a)
+  br label %if.end.icp
+
+if.end.icp:                                       ; preds = %if.false.orig_indirect, %if.true.direct_targ
+  %9 = phi i32 [ %call, %if.false.orig_indirect ], [ %call.i, %if.true.direct_targ ]
+  ret i32 %9
+}
+
+declare i1 @llvm.type.test(i8*, metadata)
+
+declare void @llvm.assume(i1)
+
+declare dso_local i32 @_ZN1B3fooEv(%class.B* %this) unnamed_addr
+declare dso_local i32 @_ZN1B3barEv(%class.B*) unnamed_addr
+declare dso_local i32 @_ZN1A3barEv(%class.A* %this) unnamed_addr
+declare dso_local i32 @_ZN1A3fooEv(%class.A* %this) unnamed_addr
+
+!0 = !{i64 16, !"_ZTS1A"}
+!1 = !{i64 16, !"_ZTS1B"}