decouple cfi from thinlto

Author: mtrofin

llvm/include/llvm/IR/ModuleSummaryIndex.h

@@ -18,11 +18,13 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/GlobalValue.h"
@@ -1317,72 +1319,75 @@ struct TypeIdSummary {
   std::map<uint64_t, WholeProgramDevirtResolution> WPDRes;
 };
 
+/// Encapsulate the names of CFI target functions. It interfaces with ThinLTO to
+/// determine efficiently which of the names need to be exported for a
+/// particular module.
 class CfiFunctionIndex {
-  DenseMap<GlobalValue::GUID, std::set<std::string, std::less<>>> Index;
-  using IndexIterator =
-      DenseMap<GlobalValue::GUID,
-               std::set<std::string, std::less<>>>::const_iterator;
-  using NestedIterator = std::set<std::string, std::less<>>::const_iterator;
+  // `Names` is the authoritative source of data. `ThinLTOToNamesIndex` is there
+  // just to efficiently retrieve which names in this index need exporting for
+  // a particular module index. We cannot guarantee the ThinLTO GUIDs are
+  // collision - free, so we associate a collection to a guid. Functions with
+  // the same name may have different GUIDs, too. So we index a list of names
+  // with the same GUID under that GUID key. We don't need the reverse because
+  // the queries from ThinLTO use GUIDs as key.
+  // Note that StringSet rehashing doesn't move keys, so we can safely store the
+  // StringRef value inserted in `Names` in ThinLTOToNamesIndex, and avoid
+  // copies.
+  // Design note: we could do away with Names and use ThinLTOToNamesIndex as
+  // index and data source, but opted against, for a small heap penalty, to
+  // avoid confusion wrt the role GUIDs play in this case: they are an artifact
+  // of the need to interface with ThinLTO, not otherwise necessary to CFI.
+  StringSet<> Names;
+
+  using InternalIndexGroup = SetVector<StringRef>;
+  DenseMap<GlobalValue::GUID, InternalIndexGroup> ThinLTOToNamesIndex;
+
+  using NestedIterator = InternalIndexGroup::const_iterator;
 
 public:
-  // Iterates keys of the DenseMap.
-  class GUIDIterator : public iterator_adaptor_base<GUIDIterator, IndexIterator,
-                                                    std::forward_iterator_tag,
-                                                    GlobalValue::GUID> {
-    using base = GUIDIterator::iterator_adaptor_base;
-
-  public:
-    GUIDIterator() = default;
-    explicit GUIDIterator(IndexIterator I) : base(I) {}
-
-    GlobalValue::GUID operator*() const { return this->wrapped()->first; }
-  };
-
   CfiFunctionIndex() = default;
-  template <typename It> CfiFunctionIndex(It B, It E) {
-    for (; B != E; ++B)
-      emplace(*B);
-  }
 
-  std::vector<StringRef> symbols() const {
-    std::vector<StringRef> Symbols;
-    for (auto &[GUID, Syms] : Index) {
-      (void)GUID;
-      llvm::append_range(Symbols, Syms);
-    }
+  /// API used for serialization, e.g. YAML.
+  std::vector<std::pair<StringRef, GlobalValue::GUID>>
+  getSortedSymbols() const {
+    std::vector<std::pair<StringRef, GlobalValue::GUID>> Symbols;
+    for (auto &[GUID, Names] : ThinLTOToNamesIndex)
+      for (auto Name : Names)
+        Symbols.emplace_back(Name, GUID);
+    llvm::sort(Symbols, std::less<>());
     return Symbols;
   }
 
-  GUIDIterator guid_begin() const { return GUIDIterator(Index.begin()); }
-  GUIDIterator guid_end() const { return GUIDIterator(Index.end()); }
-  iterator_range<GUIDIterator> guids() const {
-    return make_range(guid_begin(), guid_end());
+  /// get the set of GUIDs that should also be exported because they are the
+  /// GUIDs of the cfi functions encapsulated here.
+  auto getExportedThinLTOGUIDs() const {
+    return map_range(ThinLTOToNamesIndex, [](auto I) { return I.first; });
   }
 
-  iterator_range<NestedIterator> forGuid(GlobalValue::GUID GUID) const {
-    auto I = Index.find(GUID);
-    if (I == Index.end())
+  /// get the name(s) associated with a given ThinLTO GUID. This enables
+  /// efficient identification of the subset of names that should be included in
+  /// a module summary.
+  auto getMatchingNamesForThinLTOGUID(GlobalValue::GUID GUID) const {
+    auto I = ThinLTOToNamesIndex.find(GUID);
+    if (I == ThinLTOToNamesIndex.end())
       return make_range(NestedIterator{}, NestedIterator{});
     return make_range(I->second.begin(), I->second.end());
   }
 
-  template <typename... Args> void emplace(Args &&...A) {
-    StringRef S(std::forward<Args>(A)...);
-    GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
-        GlobalValue::dropLLVMManglingEscape(S));
-    Index[GUID].emplace(S);
+  /// Add the function name and the GUID that ThinLTO uses for it.
+  void addSymbolWithThinLTOGUID(StringRef Name, GlobalValue::GUID GUID) {
+    auto [Iter, _] = Names.insert(Name);
+    ThinLTOToNamesIndex[GUID].insert(Iter->first());
   }
 
-  size_t count(StringRef S) const {
-    GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
-        GlobalValue::dropLLVMManglingEscape(S));
-    auto I = Index.find(GUID);
-    if (I == Index.end())
-      return 0;
-    return I->second.count(S);
+  bool contains(StringRef Name) const {
+    return Names.find(Name) != Names.end();
   }
 
-  bool empty() const { return Index.empty(); }
+  bool empty() const {
+    assert(Names.empty() == ThinLTOToNamesIndex.empty());
+    return Names.empty();
+  }
 };
 
 /// 160 bits SHA1

llvm/include/llvm/IR/ModuleSummaryIndexYAML.h

@@ -327,6 +327,24 @@ template <> struct CustomMappingTraits<TypeIdSummaryMapTy> {
   }
 };
 
+template <> struct MappingTraits<std::pair<StringRef, GlobalValue::GUID>> {
+  static void mapping(IO &io,
+                      std::pair<StringRef, GlobalValue::GUID> &NameAndGUID) {
+    io.mapRequired("Name", NameAndGUID.first);
+    io.mapRequired("GUID", NameAndGUID.second);
+  }
+};
+
+using StringAndGUID = std::pair<llvm::StringRef, llvm::GlobalValue::GUID>;
+
+} // namespace yaml
+} // namespace llvm
+
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::StringAndGUID)
+
+namespace llvm {
+namespace yaml {
+
 template <> struct MappingTraits<ModuleSummaryIndex> {
   static void mapping(IO &io, ModuleSummaryIndex& index) {
     io.mapOptional("GlobalValueMap", index.GlobalValueMap);
@@ -352,25 +370,25 @@ template <> struct MappingTraits<ModuleSummaryIndex> {
                    index.WithGlobalValueDeadStripping);
 
     if (io.outputting()) {
-      auto CfiFunctionDefs = index.CfiFunctionDefs.symbols();
-      llvm::sort(CfiFunctionDefs);
+      auto CfiFunctionDefs = index.CfiFunctionDefs.getSortedSymbols();
       io.mapOptional("CfiFunctionDefs", CfiFunctionDefs);
-      auto CfiFunctionDecls(index.CfiFunctionDecls.symbols());
-      llvm::sort(CfiFunctionDecls);
+      auto CfiFunctionDecls(index.CfiFunctionDecls.getSortedSymbols());
       io.mapOptional("CfiFunctionDecls", CfiFunctionDecls);
     } else {
-      std::vector<std::string> CfiFunctionDefs;
+      std::vector<std::pair<StringRef, GlobalValue::GUID>> CfiFunctionDefs;
       io.mapOptional("CfiFunctionDefs", CfiFunctionDefs);
-      index.CfiFunctionDefs = {CfiFunctionDefs.begin(), CfiFunctionDefs.end()};
-      std::vector<std::string> CfiFunctionDecls;
+      for (auto &[S, G] : CfiFunctionDefs)
+        index.CfiFunctionDefs.addSymbolWithThinLTOGUID(S, G);
+      std::vector<std::pair<StringRef, GlobalValue::GUID>> CfiFunctionDecls;
       io.mapOptional("CfiFunctionDecls", CfiFunctionDecls);
-      index.CfiFunctionDecls = {CfiFunctionDecls.begin(),
-                                CfiFunctionDecls.end()};
+      for (auto &[S, G] : CfiFunctionDecls)
+        index.CfiFunctionDecls.addSymbolWithThinLTOGUID(S, G);
+      ;
     }
   }
 };
 
-} // End yaml namespace
-} // End llvm namespace
+} // namespace yaml
+} // namespace llvm
 
 #endif

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -8146,17 +8146,23 @@ Error ModuleSummaryIndexBitcodeReader::parseEntireSummary(unsigned ID) {
 
     case bitc::FS_CFI_FUNCTION_DEFS: {
       auto &CfiFunctionDefs = TheIndex.cfiFunctionDefs();
-      for (unsigned I = 0; I != Record.size(); I += 2)
-        CfiFunctionDefs.emplace(Strtab.data() + Record[I],
-                                static_cast<size_t>(Record[I + 1]));
+      for (unsigned I = 0; I != Record.size(); I += 3) {
+        GlobalValue::GUID ThinLTOGUID = Record[I];
+        StringRef Name(Strtab.data() + Record[I + 1],
+                       static_cast<size_t>(Record[I + 2]));
+        CfiFunctionDefs.addSymbolWithThinLTOGUID(Name, ThinLTOGUID);
+      }
       break;
     }
 
     case bitc::FS_CFI_FUNCTION_DECLS: {
       auto &CfiFunctionDecls = TheIndex.cfiFunctionDecls();
-      for (unsigned I = 0; I != Record.size(); I += 2)
-        CfiFunctionDecls.emplace(Strtab.data() + Record[I],
-                                 static_cast<size_t>(Record[I + 1]));
+      for (unsigned I = 0; I != Record.size(); I += 3) {
+        GlobalValue::GUID ThinLTOGUID = Record[I];
+        StringRef Name(Strtab.data() + Record[I + 1],
+                       static_cast<size_t>(Record[I + 2]));
+        CfiFunctionDecls.addSymbolWithThinLTOGUID(Name, ThinLTOGUID);
+      }
       break;
     }
 

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -5316,21 +5316,30 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
         getReferencedTypeIds(FS, ReferencedTypeIds);
   }
 
-  SmallVector<StringRef, 4> Functions;
+  struct CfiFunctionRecord {
+    GlobalValue::GUID GUID = 0;
+    StringRef Name;
+    bool operator<(const CfiFunctionRecord &RHS) const {
+      return Name < RHS.Name;
+    }
+  };
+  SmallVector<CfiFunctionRecord, 4> Functions;
   auto EmitCfiFunctions = [&](const CfiFunctionIndex &CfiIndex,
                               bitc::GlobalValueSummarySymtabCodes Code) {
     if (CfiIndex.empty())
       return;
     for (GlobalValue::GUID GUID : DefOrUseGUIDs) {
-      auto Defs = CfiIndex.forGuid(GUID);
-      llvm::append_range(Functions, Defs);
+      auto Names = CfiIndex.getMatchingNamesForThinLTOGUID(GUID);
+      for (StringRef Name : Names)
+        Functions.push_back({GUID, Name});
     }
     if (Functions.empty())
       return;
     llvm::sort(Functions);
-    for (const auto &S : Functions) {
-      NameVals.push_back(StrtabBuilder.add(S));
-      NameVals.push_back(S.size());
+    for (const auto &Record : Functions) {
+      NameVals.push_back(Record.GUID);
+      NameVals.push_back(StrtabBuilder.add(Record.Name));
+      NameVals.push_back(Record.Name.size());
     }
     Stream.EmitRecord(Code, NameVals);
     NameVals.clear();

llvm/lib/LTO/LTO.cpp

@@ -1580,9 +1580,9 @@ class CGThinBackend : public ThinBackendProc {
                         OnWrite, ShouldEmitImportsFiles, ThinLTOParallelism),
         ShouldEmitIndexFiles(ShouldEmitIndexFiles) {
     auto &Defs = CombinedIndex.cfiFunctionDefs();
-    CfiFunctionDefs.insert_range(Defs.guids());
+    CfiFunctionDefs.insert_range(Defs.getExportedThinLTOGUIDs());
     auto &Decls = CombinedIndex.cfiFunctionDecls();
-    CfiFunctionDecls.insert_range(Decls.guids());
+    CfiFunctionDecls.insert_range(Decls.getExportedThinLTOGUIDs());
   }
 };
 
@@ -2156,9 +2156,11 @@ Error LTO::runThinLTO(AddStreamFn AddStream, FileCache Cache,
   // Any functions referenced by the jump table in the regular LTO object must
   // be exported.
   auto &Defs = ThinLTO.CombinedIndex.cfiFunctionDefs();
-  ExportedGUIDs.insert(Defs.guid_begin(), Defs.guid_end());
+  ExportedGUIDs.insert(Defs.getExportedThinLTOGUIDs().begin(),
+                       Defs.getExportedThinLTOGUIDs().end());
   auto &Decls = ThinLTO.CombinedIndex.cfiFunctionDecls();
-  ExportedGUIDs.insert(Decls.guid_begin(), Decls.guid_end());
+  ExportedGUIDs.insert(Decls.getExportedThinLTOGUIDs().begin(),
+                       Decls.getExportedThinLTOGUIDs().end());
 
   auto isExported = [&](StringRef ModuleIdentifier, ValueInfo VI) {
     const auto &ExportList = ExportLists.find(ModuleIdentifier);

llvm/lib/Transforms/IPO/CrossDSOCFI.cpp

@@ -73,8 +73,9 @@ void CrossDSOCFI::buildCFICheck(Module &M) {
   NamedMDNode *CfiFunctionsMD = M.getNamedMetadata("cfi.functions");
   if (CfiFunctionsMD) {
     for (auto *Func : CfiFunctionsMD->operands()) {
-      assert(Func->getNumOperands() >= 2);
-      for (unsigned I = 2; I < Func->getNumOperands(); ++I)
+      assert(Func->getNumOperands() >= 3);
+      assert(isa<ConstantAsMetadata>(Func->getOperand(2)));
+      for (unsigned I = 3; I < Func->getNumOperands(); ++I)
         if (ConstantInt *TypeId =
                 extractNumericTypeId(cast<MDNode>(Func->getOperand(I).get())))
           TypeIds.insert(TypeId->getZExtValue());

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

@@ -1788,10 +1788,15 @@ void LowerTypeTestsModule::buildBitSetsFromFunctionsNative(
     }
 
     if (IsExported) {
+      // TODO: use F->getGUID() once #184065 is relanded.
+      GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
+          GlobalValue::dropLLVMManglingEscape(F->getName()));
       if (IsJumpTableCanonical)
-        ExportSummary->cfiFunctionDefs().emplace(F->getName());
+        ExportSummary->cfiFunctionDefs().addSymbolWithThinLTOGUID(F->getName(),
+                                                                  GUID);
       else
-        ExportSummary->cfiFunctionDecls().emplace(F->getName());
+        ExportSummary->cfiFunctionDecls().addSymbolWithThinLTOGUID(F->getName(),
+                                                                   GUID);
     }
 
     if (!IsJumpTableCanonical) {
@@ -2123,9 +2128,9 @@ bool LowerTypeTestsModule::lower() {
       // have the same name, but it's not the one we are looking for.
       if (F.hasLocalLinkage())
         continue;
-      if (ImportSummary->cfiFunctionDefs().count(F.getName()))
+      if (ImportSummary->cfiFunctionDefs().contains(F.getName()))
         Defs.push_back(&F);
-      else if (ImportSummary->cfiFunctionDecls().count(F.getName()))
+      else if (ImportSummary->cfiFunctionDecls().contains(F.getName()))
         Decls.push_back(&F);
     }
 
@@ -2208,8 +2213,10 @@ bool LowerTypeTestsModule::lower() {
                 ->getUniqueInteger()
                 .getZExtValue());
         const GlobalValue::GUID GUID =
-            GlobalValue::getGUIDAssumingExternalLinkage(
-                GlobalValue::dropLLVMManglingEscape(FunctionName));
+            cast<ConstantAsMetadata>(FuncMD->getOperand(2))
+                ->getValue()
+                ->getUniqueInteger()
+                .getZExtValue();
         // Do not emit jumptable entries for functions that are not-live and
         // have no live references (and are not exported with cross-DSO CFI.)
         if (!ExportSummary->isGUIDLive(GUID))
@@ -2280,7 +2287,7 @@ bool LowerTypeTestsModule::lower() {
             F->setLinkage(GlobalValue::ExternalWeakLinkage);
 
           F->eraseMetadata(LLVMContext::MD_type);
-          for (unsigned I = 2; I < FuncMD->getNumOperands(); ++I)
+          for (unsigned I = 3; I < FuncMD->getNumOperands(); ++I)
             F->addMetadata(LLVMContext::MD_type,
                            *cast<MDNode>(FuncMD->getOperand(I).get()));
         }

llvm/lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp

@@ -436,6 +436,11 @@ void splitAndWriteThinLTOBitcode(
       Linkage = CFL_Declaration;
     Elts.push_back(ConstantAsMetadata::get(
         llvm::ConstantInt::get(Type::getInt8Ty(Ctx), Linkage)));
+    // TODO: use F->getGUID() once #184065 is relanded.
+    GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
+        GlobalValue::dropLLVMManglingEscape(V->getName()));
+    Elts.push_back(ConstantAsMetadata::get(
+        llvm::ConstantInt::get(Type::getInt64Ty(Ctx), GUID)));
     append_range(Elts, Types);
     CfiFunctionMDs.push_back(MDTuple::get(Ctx, Elts));
   }

llvm/test/ThinLTO/X86/cfi-icall-only-defuse.ll

@@ -40,15 +40,15 @@ define i8 @f(i1 %i, ptr %p) {
 !0 = !{i64 0, !"t1"}
 
 ; FOOBAZ:   <GLOBALVAL_SUMMARY_BLOCK
-; FOOBAZ:     <CFI_FUNCTION_DEFS op0=0 op1=3 op2=3 op3=3 op4=6 op5=3/>
+; FOOBAZ:     <CFI_FUNCTION_DEFS op0={{.*}} op1=0 op2=3 op3={{.*}} op4=3 op5=3 op6={{.*}} op7=6 op8=3/>
 ; FOOBAZ:     <TYPE_ID op0=9 op1=2 op2=4 op3=7 op4=0 op5=0 op6=0 op7=0/>
 ; FOOBAZ:   </GLOBALVAL_SUMMARY_BLOCK>
 ; FOOBAZ:      <STRTAB_BLOCK
 ; FOOBAZ-NEXT:   <BLOB abbrevid=4/> blob data = 'barbazfoot1'
 ; FOOBAZ-NEXT: </STRTAB_BLOCK>
 
 ; BARQUX:   <GLOBALVAL_SUMMARY_BLOCK
-; BARQUX:     <CFI_FUNCTION_DEFS op0=0 op1=3 op2=3 op3=3 op4=6 op5=3/>
+; BARQUX:     <CFI_FUNCTION_DEFS op0={{.*}} op1=0 op2=3 op3={{.*}} op4=3 op5=3 op6={{.*}} op7=6 op8=3/>
 ; BARQUX:     <TYPE_ID op0=9 op1=2 op2=4 op3=7 op4=0 op5=0 op6=0 op7=0/>
 ; BARQUX:   </GLOBALVAL_SUMMARY_BLOCK>
 ; BARQUX:      <STRTAB_BLOCK