[ThinLTO][Bitcode] Generate import type in bitcode

[mingmingl-llvm]

For distributed ThinLTO, the LTO indexing step generates combined summary for each module, and postlink pipeline reads the combined summary which stores the information for link-time optimization.

This patch populates the 'import type' of a summary in bitcode, and updates bitcode reader to parse the bit correctly.

[llvmbot]

@llvm/pr-subscribers-lto
@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-ir

Author: Mingming Liu (minglotus-6)

Changes

Given a list of modules and a function f, there are multiple combinations of imports (each module may need f as definition, declaration, or doesn't reference it), and GVFlags.ImportType should be a function of <module, value-summary>.

When writing the per-module combined summary, bitcode writer looks up a map to write to the ImportType bit in the given module.

Tests:

• Regression test which depends on this change in a subsequent patch.

This is a follow-up patch of #87597.

---

Full diff: https://github.com/llvm/llvm-project/pull/87600.diff

3 Files Affected:

• (modified) llvm/include/llvm/Bitcode/BitcodeWriter.h (+10-5)
• (modified) llvm/include/llvm/IR/ModuleSummaryIndex.h (+6)
• (modified) llvm/lib/Bitcode/Writer/BitcodeWriter.cpp (+39-11)

diff --git a/llvm/include/llvm/Bitcode/BitcodeWriter.h b/llvm/include/llvm/Bitcode/BitcodeWriter.h
index 248d33f4502ef0..e94b074d9788e6 100644
--- a/llvm/include/llvm/Bitcode/BitcodeWriter.h
+++ b/llvm/include/llvm/Bitcode/BitcodeWriter.h
@@ -102,7 +102,8 @@ class raw_ostream;
 
     void writeIndex(
         const ModuleSummaryIndex *Index,
-        const std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex);
+        const std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex,
+        const ModuleToGVSummaryPtrSet *ModuleToDecSummaries);
   };
 
   /// Write the specified module to the specified raw output stream.
@@ -147,10 +148,14 @@ class raw_ostream;
   /// where it will be written in a new bitcode block. This is used when
   /// writing the combined index file for ThinLTO. When writing a subset of the
   /// index for a distributed backend, provide the \p ModuleToSummariesForIndex
-  /// map.
-  void writeIndexToFile(const ModuleSummaryIndex &Index, raw_ostream &Out,
-                        const std::map<std::string, GVSummaryMapTy>
-                            *ModuleToSummariesForIndex = nullptr);
+  /// map. For each module, \p ModuleToDecSummaries specifies the set of
+  /// summaries for which the corresponding value should be imported as a
+  /// declaration (prototype).
+  void writeIndexToFile(
+      const ModuleSummaryIndex &Index, raw_ostream &Out,
+      const std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex =
+          nullptr,
+      const ModuleToGVSummaryPtrSet *ModuleToDecSummaries = nullptr);
 
   /// If EmbedBitcode is set, save a copy of the llvm IR as data in the
   ///  __LLVM,__bitcode section (.llvmbc on non-MacOS).
diff --git a/llvm/include/llvm/IR/ModuleSummaryIndex.h b/llvm/include/llvm/IR/ModuleSummaryIndex.h
index 5d137d4b3553cf..66209b8cf3ea2d 100644
--- a/llvm/include/llvm/IR/ModuleSummaryIndex.h
+++ b/llvm/include/llvm/IR/ModuleSummaryIndex.h
@@ -1272,6 +1272,12 @@ using ModulePathStringTableTy = StringMap<ModuleHash>;
 /// a particular module, and provide efficient access to their summary.
 using GVSummaryMapTy = DenseMap<GlobalValue::GUID, GlobalValueSummary *>;
 
+/// A set of global value summary pointers.
+using GVSummaryPtrSet = SmallPtrSet<GlobalValueSummary *, 4>;
+
+/// The key is module path, and value is a set of global value summary pointers.
+using ModuleToGVSummaryPtrSet = std::map<std::string, GVSummaryPtrSet>;
+
 /// Map of a type GUID to type id string and summary (multimap used
 /// in case of GUID conflicts).
 using TypeIdSummaryMapTy =
diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
index 6d01e3b4d82189..a060265a9bc8dd 100644
--- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -428,6 +428,10 @@ class IndexBitcodeWriter : public BitcodeWriterBase {
   /// The combined index to write to bitcode.
   const ModuleSummaryIndex &Index;
 
+  /// For each module, provides the set of global value summaries for which the
+  /// value (function, function alias, etc) should be imported as a declaration.
+  const ModuleToGVSummaryPtrSet *ModuleToDecSummaries = nullptr;
+
   /// When writing a subset of the index for distributed backends, client
   /// provides a map of modules to the corresponding GUIDs/summaries to write.
   const std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex;
@@ -452,11 +456,17 @@ class IndexBitcodeWriter : public BitcodeWriterBase {
   /// Constructs a IndexBitcodeWriter object for the given combined index,
   /// writing to the provided \p Buffer. When writing a subset of the index
   /// for a distributed backend, provide a \p ModuleToSummariesForIndex map.
-  IndexBitcodeWriter(BitstreamWriter &Stream, StringTableBuilder &StrtabBuilder,
-                     const ModuleSummaryIndex &Index,
-                     const std::map<std::string, GVSummaryMapTy>
-                         *ModuleToSummariesForIndex = nullptr)
+  /// If provided, \p ModuleToDecSummaries specifies the set of summaries for
+  /// which the corresponding functions or aliased functions should be imported
+  /// as a declaration (but not definition) for each module.
+  IndexBitcodeWriter(
+      BitstreamWriter &Stream, StringTableBuilder &StrtabBuilder,
+      const ModuleSummaryIndex &Index,
+      const ModuleToGVSummaryPtrSet *ModuleToDecSummaries = nullptr,
+      const std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex =
+          nullptr)
       : BitcodeWriterBase(Stream, StrtabBuilder), Index(Index),
+        ModuleToDecSummaries(ModuleToDecSummaries),
         ModuleToSummariesForIndex(ModuleToSummariesForIndex) {
     // Assign unique value ids to all summaries to be written, for use
     // in writing out the call graph edges. Save the mapping from GUID
@@ -1202,7 +1212,8 @@ static uint64_t getEncodedFFlags(FunctionSummary::FFlags Flags) {
 
 // Decode the flags for GlobalValue in the summary. See getDecodedGVSummaryFlags
 // in BitcodeReader.cpp.
-static uint64_t getEncodedGVSummaryFlags(GlobalValueSummary::GVFlags Flags) {
+static uint64_t getEncodedGVSummaryFlags(GlobalValueSummary::GVFlags Flags,
+                                         bool ImportAsDecl = false) {
   uint64_t RawFlags = 0;
 
   RawFlags |= Flags.NotEligibleToImport; // bool
@@ -1217,7 +1228,8 @@ static uint64_t getEncodedGVSummaryFlags(GlobalValueSummary::GVFlags Flags) {
 
   RawFlags |= (Flags.Visibility << 8); // 2 bits
 
-  RawFlags |= (Flags.ImportType << 10); // 1 bit
+  unsigned ImportType = Flags.ImportType | ImportAsDecl;
+  RawFlags |= (ImportType << 10); // 1 bit
 
   return RawFlags;
 }
@@ -4543,6 +4555,17 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
   Abbv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::VBR, 8));
   unsigned AllocAbbrev = Stream.EmitAbbrev(std::move(Abbv));
 
+  auto shouldImportValueAsDecl = [&](GlobalValueSummary *GVS) -> bool {
+    if (ModuleToDecSummaries == nullptr)
+      return false;
+    auto Iter = ModuleToDecSummaries->find(GVS->modulePath().str());
+    if (Iter == ModuleToDecSummaries->end())
+      return false;
+    // For the current module, the value for GVS should be imported as a
+    // declaration.
+    return Iter->second.contains(GVS);
+  };
+
   // The aliases are emitted as a post-pass, and will point to the value
   // id of the aliasee. Save them in a vector for post-processing.
   SmallVector<AliasSummary *, 64> Aliases;
@@ -4653,7 +4676,8 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
     NameVals.push_back(*ValueId);
     assert(ModuleIdMap.count(FS->modulePath()));
     NameVals.push_back(ModuleIdMap[FS->modulePath()]);
-    NameVals.push_back(getEncodedGVSummaryFlags(FS->flags()));
+    NameVals.push_back(
+        getEncodedGVSummaryFlags(FS->flags(), shouldImportValueAsDecl(FS)));
     NameVals.push_back(FS->instCount());
     NameVals.push_back(getEncodedFFlags(FS->fflags()));
     NameVals.push_back(FS->entryCount());
@@ -4702,7 +4726,8 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
     NameVals.push_back(AliasValueId);
     assert(ModuleIdMap.count(AS->modulePath()));
     NameVals.push_back(ModuleIdMap[AS->modulePath()]);
-    NameVals.push_back(getEncodedGVSummaryFlags(AS->flags()));
+    NameVals.push_back(
+        getEncodedGVSummaryFlags(AS->flags(), shouldImportValueAsDecl(AS)));
     auto AliaseeValueId = SummaryToValueIdMap[&AS->getAliasee()];
     assert(AliaseeValueId);
     NameVals.push_back(AliaseeValueId);
@@ -5036,8 +5061,10 @@ void BitcodeWriter::writeModule(const Module &M,
 
 void BitcodeWriter::writeIndex(
     const ModuleSummaryIndex *Index,
-    const std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex) {
+    const std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex,
+    const ModuleToGVSummaryPtrSet *ModuleToDecSummaries) {
   IndexBitcodeWriter IndexWriter(*Stream, StrtabBuilder, *Index,
+                                 ModuleToDecSummaries,
                                  ModuleToSummariesForIndex);
   IndexWriter.write();
 }
@@ -5090,12 +5117,13 @@ void IndexBitcodeWriter::write() {
 // index for a distributed backend, provide a \p ModuleToSummariesForIndex map.
 void llvm::writeIndexToFile(
     const ModuleSummaryIndex &Index, raw_ostream &Out,
-    const std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex) {
+    const std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex,
+    const ModuleToGVSummaryPtrSet *ModuleToDecSummaries) {
   SmallVector<char, 0> Buffer;
   Buffer.reserve(256 * 1024);
 
   BitcodeWriter Writer(Buffer);
-  Writer.writeIndex(&Index, ModuleToSummariesForIndex);
+  Writer.writeIndex(&Index, ModuleToSummariesForIndex, ModuleToDecSummaries);
   Writer.writeStrtab();
 
   Out.write((char *)&Buffer.front(), Buffer.size());

[mingmingl-llvm]

Tagging @jvoung for review since I cannot add you, thanks!

[jvoung]

Not a big issue, but was wondering if the module path is needed or if a set instead of map would work? It seemed like each GlobalValueSummary* would only be imported from one module for declaration. Or is it possible that for a particular GVS_i ?

mod_A -> GVS_i
mod_B -> GVS_i

Otherwise, the map does match the gatherImportedSummariesForModule more closely.

[mingmingl-llvm]

Good question. I think it makes sense to use a set of pointers as suggested, given ModuleToGVSummaryPtrSet is computed using pointers as opposed to GUIDs. Going to make the change.

While it's possible to have {mod_A, GUID_foo} and {mod_B, GUID_foo} entries because of comdats, each GlobalValueSummary* will be associated with exactly one module.

[mingmingl-llvm]

Done. Going to update the subsequent patch after updating this PR.

[jvoung]

Ah got it, thanks for the clarification!

[jvoung]

nit: maybe indicate "When writing a subset of the index" also?

[mingmingl-llvm]

Updated the comment to "When writing combined summaries, provides the set of global value summaries...".

While we write a subset of the index only for combined summaries, 'combined summary' is slightly preferred as it hints this map doesn't apply for per-module summary or the index.