[ThinLTO][Bitcode] Generate import type in bitcode

[minglotus-6]

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());

[minglotus-6]

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.

[minglotus-6]

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.

[minglotus-6]

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?

[minglotus-6]

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.

[jvoung]

Can you add a documentation for the DecSummaries parameter here too? (similar to the ThinLTOCodeGenerator.h comment, might be useful to indicate it's a subset of ModuleToSummariesForIndex).

[minglotus-6]

done.

[jvoung]

Nice!

[jvoung]

It might be interesting to test happens if "calleeAddrs" had just an alias of a function and not the original function (right now both large_indirect_callee and large_indirect_callee_alias are there?), but could be okay to test that later too.

[minglotus-6]

I updated the test case as suggested, and it's indeed an interesting case, and actually needs a follow-up fix. As shown by the updated test, function alias summary has 'null' aliasee. I'm preparing a fix.

For function alias definition, postlink pipeline clones the aliasee function and keeps using alias's name. I'm looking into whether it makes sense to rewrite the callee to aliasee (for indirect calls we need to update value profiles).

[minglotus-6]

As shown by the updated test, function alias summary has 'null' aliasee. I'm preparing a fix.

Update on this: it's valid for an aliasee to be 'null'.

llvm-project/llvm/test/ThinLTO/X86/distributed_indexes.ll

Lines 45 to 49 in 3e15c97

	; Make sure that when the alias is imported as a copy of the aliasee, but the
	; aliasee is not being imported by itself, that we can still print the summary.
	; The aliasee should be "null".
	; RUN: llvm-dis %t1.bc.thinlto.bc -o - | FileCheck %s --check-prefix=DIS
	; DIS: aliasee: null

is a test case

[jvoung]

Thanks for checking!