[CFI][ThinLTO] Remove the need for CFI calculating ThinLTO GUIDs

llvm/include/llvm/IR/AutoUpgrade.h

@@ -63,6 +63,10 @@ namespace llvm {
   /// module is modified.
   LLVM_ABI bool UpgradeModuleFlags(Module &M);
 
+  /// Upgrade the cfi.functions metadata node by calculating and inserting
+  /// the GUID for each function entry if it's missing.
+  LLVM_ABI bool UpgradeCFIFunctionsMetadata(Module &M);
+
   /// Convert legacy nvvm.annotations metadata to appropriate function
   /// attributes.
   LLVM_ABI void UpgradeNVVMAnnotations(Module &M);

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
@@ -1566,7 +1571,7 @@ class ModuleSummaryIndex {
   // in the way some record are interpreted, like flags for instance.
   // Note that incrementing this may require changes in both BitcodeReader.cpp
   // and BitcodeWriter.cpp.
-  static constexpr uint64_t BitcodeSummaryVersion = 13;
+  static constexpr uint64_t BitcodeSummaryVersion = 14;
 
   // Regular LTO module name for ASM writer
   static constexpr const char *getRegularLTOModuleName() {

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/AsmParser/LLParser.cpp

@@ -476,6 +476,7 @@ bool LLParser::validateEndOfModule(bool UpgradeDebugInfo) {
     llvm::UpgradeDebugInfo(*M);
 
   UpgradeModuleFlags(*M);
+  UpgradeCFIFunctionsMetadata(*M);
   UpgradeNVVMAnnotations(*M);
   UpgradeSectionAttributes(*M);
   copyModuleAttrToFunctions(*M);

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -3979,6 +3979,8 @@ Error BitcodeReader::materializeMetadata() {
     }
   }
 
+  UpgradeCFIFunctionsMetadata(*TheModule);
+
   DeferredMetadataInfo.clear();
   return Error::success();
 }
@@ -8146,17 +8148,43 @@ 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]));
+      if (Version < 14) {
+        for (unsigned I = 0; I != Record.size(); I += 2) {
+          StringRef Name(Strtab.data() + Record[I],
+                         static_cast<size_t>(Record[I + 1]));
+          GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
+              GlobalValue::dropLLVMManglingEscape(Name));
+          CfiFunctionDefs.addSymbolWithThinLTOGUID(Name, GUID);
+        }
+      } else {
+        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]));
+      if (Version < 14) {
+        for (unsigned I = 0; I != Record.size(); I += 2) {
+          StringRef Name(Strtab.data() + Record[I],
+                         static_cast<size_t>(Record[I + 1]));
+          GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
+              GlobalValue::dropLLVMManglingEscape(Name));
+          CfiFunctionDecls.addSymbolWithThinLTOGUID(Name, GUID);
+        }
+      } else {
+        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/IR/AutoUpgrade.cpp

@@ -26,6 +26,7 @@
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/Instruction.h"
@@ -6359,6 +6360,46 @@ bool llvm::UpgradeModuleFlags(Module &M) {
   return Changed;
 }
 
+bool llvm::UpgradeCFIFunctionsMetadata(Module &M) {
+  NamedMDNode *CFIConsts = M.getNamedMetadata("cfi.functions");
+  if (!CFIConsts)
+    return false;
+
+  bool Changed = false;
+  for (unsigned I = 0, E = CFIConsts->getNumOperands(); I != E; ++I) {
+    MDNode *Op = CFIConsts->getOperand(I);
+    if (Op->getNumOperands() >= 3 &&
+        isa<ConstantAsMetadata>(Op->getOperand(2)) &&
+        cast<ConstantAsMetadata>(Op->getOperand(2))->getType()->isIntegerTy(64))
+      continue;
+
+    if (Op->getNumOperands() < 2)
+      continue;
+
+    MDString *NameMD = dyn_cast_or_null<MDString>(Op->getOperand(0));
+    if (!NameMD)
+      continue;
+
+    StringRef Name = NameMD->getString();
+    GlobalValue::GUID GUID = GlobalValue::getGUIDAssumingExternalLinkage(
+        GlobalValue::dropLLVMManglingEscape(Name));
+
+    SmallVector<Metadata *, 4> Elts;
+    Elts.push_back(Op->getOperand(0));
+    Elts.push_back(Op->getOperand(1));
+    Elts.push_back(ConstantAsMetadata::get(
+        ConstantInt::get(Type::getInt64Ty(M.getContext()), GUID)));
+
+    for (unsigned J = 2, EJ = Op->getNumOperands(); J != EJ; ++J)
+      Elts.push_back(Op->getOperand(J));
+
+    CFIConsts->setOperand(I, MDNode::get(M.getContext(), Elts));
+    Changed = true;
+  }
+
+  return Changed;
+}
+
 void llvm::UpgradeSectionAttributes(Module &M) {
   auto TrimSpaces = [](StringRef Section) -> std::string {
     SmallVector<StringRef, 5> Components;