[ThinLTO/WPD] Ensure devirtualized targets use promoted symbol when necessary

[tejohnson]

Summary

This fixes a hole in the handling of devirtualized targets that were
local but need promoting due to devirtualization in another module. We
were not correctly referencing the promoted symbol in some cases. Make
sure the code that updates the name also looks at the ExportedGUIDs set
by utilizing a callback that checks all conditions (the callback
utilized by the internalization/promotion code).

[hiraditya]

lib/LTO/LTO.cpp

1335

Maybe we could make this a separate function, it seems we are using similar checks in llvm/lib/LTO/ThinLTOCodeGenerator.cpp:528,

[tejohnson]

lib/LTO/LTO.cpp

1335

Do you mean the isExported lambda (ThinLTOCodeGenerator.cpp:528 is unrelated code at HEAD)? We use a call back in part because the different LTO APIs handle exported symbols differently. The ExportedGUIDs set here in the new API is set up differently than the GUIDPreservedSymbols in the old API. Also, the old LTO API doesn't currently do WPD.

[hiraditya]

lib/LTO/LTO.cpp

1335

oh sorry, i was referring to: https://github.com/llvm-mirror/llvm/blob/master/lib/LTO/ThinLTOCodeGenerator.cpp#L586. But i understood the point. thanks!

include/llvm/Transforms/IPO/WholeProgramDevirt.h

 void runWholeProgramDevirtOnIndex(
     ModuleSummaryIndex &Summary, std::set<GlobalValue::GUID> &ExportedGUIDs,
     std::map<ValueInfo, std::vector<VTableSlotSummary>> &LocalWPDTargetsMap);
 
 /// Call after cross-module importing to update the recorded single impl
 /// devirt target names for any locals that were exported.
 void updateIndexWPDForExports(
     ModuleSummaryIndex &Summary,
-    StringMap<FunctionImporter::ExportSetTy> &ExportLists,
+    function_ref<bool(StringRef, GlobalValue::GUID)> isExported,
     std::map<ValueInfo, std::vector<VTableSlotSummary>> &LocalWPDTargetsMap);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_IPO_WHOLEPROGRAMDEVIRT_H

lib/LTO/LTO.cpp

Error LTO::runThinLTO(AddStreamFn AddStream, NativeObjectCache Cache,
   std::map<ValueInfo, std::vector<VTableSlotSummary>> LocalWPDTargetsMap;
   runWholeProgramDevirtOnIndex(ThinLTO.CombinedIndex, ExportedGUIDs,
                                LocalWPDTargetsMap);
 
   if (Conf.OptLevel > 0)
     ComputeCrossModuleImport(ThinLTO.CombinedIndex, ModuleToDefinedGVSummaries,
                              ImportLists, ExportLists);
 
-  // Update local devirtualized targets that were exported by cross-module
-  // importing
-  updateIndexWPDForExports(ThinLTO.CombinedIndex, ExportLists,
-                           LocalWPDTargetsMap);
-
   // Figure out which symbols need to be internalized. This also needs to happen
   // at -O0 because summary-based DCE is implemented using internalization, and
   // we must apply DCE consistently with the full LTO module in order to avoid
   // undefined references during the final link.
   for (auto &Res : GlobalResolutions) {
     // If the symbol does not have external references or it is not prevailing,
     // then not need to mark it as exported from a ThinLTO partition.
     if (Res.second.Partition != GlobalResolution::External ||
for (auto &Def : ThinLTO.CombinedIndex.cfiFunctionDefs())
     ExportedGUIDs.insert(
         GlobalValue::getGUID(GlobalValue::dropLLVMManglingEscape(Def)));
 
   auto isExported = [&](StringRef ModuleIdentifier, GlobalValue::GUID GUID) {
     const auto &ExportList = ExportLists.find(ModuleIdentifier);
     return (ExportList != ExportLists.end() &&
             ExportList->second.count(GUID)) ||
            ExportedGUIDs.count(GUID);
   };
+
+  // Update local devirtualized targets that were exported by cross-module
+  // importing or by other devirtualizations marked in the ExportedGUIDs set.
+  updateIndexWPDForExports(ThinLTO.CombinedIndex, isExported,
+                           LocalWPDTargetsMap);
+
   auto isPrevailing = [&](GlobalValue::GUID GUID,
                           const GlobalValueSummary *S) {
     return ThinLTO.PrevailingModuleForGUID[GUID] == S->modulePath();
   };
   thinLTOInternalizeAndPromoteInIndex(ThinLTO.CombinedIndex, isExported,
                                       isPrevailing);
 
   auto recordNewLinkage = [&](StringRef ModuleIdentifier,

lib/Transforms/IPO/WholeProgramDevirt.cpp

 void runWholeProgramDevirtOnIndex(
     ModuleSummaryIndex &Summary, std::set<GlobalValue::GUID> &ExportedGUIDs,
     std::map<ValueInfo, std::vector<VTableSlotSummary>> &LocalWPDTargetsMap) {
   DevirtIndex(Summary, ExportedGUIDs, LocalWPDTargetsMap).run();
 }
 
 void updateIndexWPDForExports(
     ModuleSummaryIndex &Summary,
-    StringMap<FunctionImporter::ExportSetTy> &ExportLists,
+    function_ref<bool(StringRef, GlobalValue::GUID)> isExported,
     std::map<ValueInfo, std::vector<VTableSlotSummary>> &LocalWPDTargetsMap) {
   for (auto &T : LocalWPDTargetsMap) {
     auto &VI = T.first;
     // This was enforced earlier during trySingleImplDevirt.
     assert(VI.getSummaryList().size() == 1 &&
            "Devirt of local target has more than one copy");
     auto &S = VI.getSummaryList()[0];
-    const auto &ExportList = ExportLists.find(S->modulePath());
-    if (ExportList == ExportLists.end() ||
-        !ExportList->second.count(VI.getGUID()))
+    if (!isExported(S->modulePath(), VI.getGUID()))
       continue;
 
     // It's been exported by a cross module import.
     for (auto &SlotSummary : T.second) {
       auto *TIdSum = Summary.getTypeIdSummary(SlotSummary.TypeID);
       assert(TIdSum);
       auto WPDRes = TIdSum->WPDRes.find(SlotSummary.ByteOffset);
       assert(WPDRes != TIdSum->WPDRes.end());

test/ThinLTO/X86/Inputs/devirt_promote.ll

+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+%struct.A = type { i32 (...)** }
+%struct.B = type { %struct.A }
+
+@_ZTV1B = constant { [4 x i8*] } { [4 x i8*] [i8* null, i8* undef, i8* bitcast (i32 (%struct.B*, i32)* @_ZN1B1fEi to i8*), i8* bitcast (i32 (%struct.A*, i32)* @_ZN1A1nEi to i8*)] }, !type !0, !type !1
+
+define i32 @_ZN1B1fEi(%struct.B* %this, i32 %a) #0 {
+   ret i32 0;
+}
+
+define internal i32 @_ZN1A1nEi(%struct.A* %this, i32 %a) #0 {
+   ret i32 0;
+}
+
+define i32 @test2(%struct.B* %obj, i32 %a) {
+entry:
+  %0 = bitcast %struct.B* %obj to i8***
+  %vtable2 = load i8**, i8*** %0
+  %1 = bitcast i8** %vtable2 to i8*
+  %p2 = call i1 @llvm.type.test(i8* %1, metadata !"_ZTS1B")
+  call void @llvm.assume(i1 %p2)
+
+  %fptrptr = getelementptr i8*, i8** %vtable2, i32 1
+  %2 = bitcast i8** %fptrptr to i32 (%struct.B*, i32)**
+  %fptr33 = load i32 (%struct.B*, i32)*, i32 (%struct.B*, i32)** %2, align 8
+
+  %call4 = tail call i32 %fptr33(%struct.B* nonnull %obj, i32 %a)
+  ret i32 %call4
+}
+
+attributes #0 = { noinline optnone }
+
+declare i1 @llvm.type.test(i8*, metadata)
+declare void @llvm.assume(i1)
+
+!0 = !{i64 16, !"_ZTS1A"}
+!1 = !{i64 16, !"_ZTS1B"}

test/ThinLTO/X86/devirt_promote.ll

+; REQUIRES: x86-registered-target
+
+; Test devirtualization requiring promotion of local targets, where the
+; promotion is required by one devirtualization and needs to be updated
+; for a second devirtualization in the defining module as a post-pass
+; update.
+
+; Generate unsplit module with summary for ThinLTO index-based WPD.
+; RUN: opt -thinlto-bc -o %t3.o %s
+; RUN: opt -thinlto-bc -o %t4.o %p/Inputs/devirt_promote.ll
+
+; RUN: llvm-lto2 run %t3.o %t4.o -save-temps -use-new-pm -pass-remarks=. \
+; RUN:   -wholeprogramdevirt-print-index-based \
+; RUN:   -o %t5 \
+; RUN:   -r=%t3.o,test,px \
+; RUN:   -r=%t4.o,_ZN1B1fEi,p \
+; RUN:   -r=%t4.o,test2,px \
+; RUN:   -r=%t4.o,_ZTV1B,px \
+; RUN:   2>&1 | FileCheck %s --check-prefix=REMARK --check-prefix=PRINT
+; RUN: llvm-dis %t5.1.4.opt.bc -o - | FileCheck %s --check-prefix=CHECK-IR1
+; RUN: llvm-dis %t5.2.4.opt.bc -o - | FileCheck %s --check-prefix=CHECK-IR2
+; RUN: llvm-nm %t5.1 | FileCheck %s --check-prefix=NM-INDEX1
+; RUN: llvm-nm %t5.2 | FileCheck %s --check-prefix=NM-INDEX2
+
+; NM-INDEX1: U _ZN1A1nEi.llvm.
+
+; Make sure that not only did _ZN1A1nEi get promoted (due to the
+; devirtualization in the other module) but the reference due to the
+; devirtualization in its defining module should be to the promoted
+; symbol.
+; NM-INDEX2-NOT: U _ZN1A1nEi
+; NM-INDEX2: T _ZN1A1nEi.llvm.
+; NM-INDEX2-NOT: U _ZN1A1nEi
+
+; We should devirt call to _ZN1A1nEi once in importing module and once
+; in original (exporting) module.
+; REMARK-COUNT-2: single-impl: devirtualized a call to _ZN1A1nEi.llvm.
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+%struct.A = type { i32 (...)** }
+
+; CHECK-IR1-LABEL: define i32 @test
+define i32 @test(%struct.A* %obj, i32 %a) {
+entry:
+  %0 = bitcast %struct.A* %obj to i8***
+  %vtable = load i8**, i8*** %0
+  %1 = bitcast i8** %vtable to i8*
+  %p = call i1 @llvm.type.test(i8* %1, metadata !"_ZTS1A")
+  call void @llvm.assume(i1 %p)
+  %fptrptr = getelementptr i8*, i8** %vtable, i32 1
+  %2 = bitcast i8** %fptrptr to i32 (%struct.A*, i32)**
+  %fptr1 = load i32 (%struct.A*, i32)*, i32 (%struct.A*, i32)** %2, align 8
+
+  ; Check that the call was devirtualized.
+  ; CHECK-IR1: %call = tail call i32 bitcast (void ()* @_ZN1A1nEi
+  %call = tail call i32 %fptr1(%struct.A* nonnull %obj, i32 %a)
+
+  ret i32 %call
+}
+; CHECK-IR1-LABEL: ret i32
+; CHECK-IR1-LABEL: }
+
+; CHECK-IR2: define i32 @test2
+; Check that the call was devirtualized.
+; CHECK-IR2:   %call4 = tail call i32 @_ZN1A1nEi
+
+declare i1 @llvm.type.test(i8*, metadata)
+declare void @llvm.assume(i1)
+
+attributes #0 = { noinline optnone }