[ThinLTO] Emit individual index files for distributed backends

Summary:
When launching ThinLTO backends in a distributed build (currently
supported in gold via the thinlto-index-only plugin option), emit
an individual index file for each backend process as described here:
http://lists.llvm.org/pipermail/llvm-dev/2016-April/098272.html

The individual index file encodes the summary and module information
required for implementing the importing/exporting decisions made
for a given module in the thin link step.
This is in place of the current mechanism that uses the combined index
to make importing decisions in each back end independently. It is an
enabler for doing global summary based optimizations in the thin link
step (which will be recorded in the individual index files), and reduces
the size of the index that must be sent to each backend process, and
the amount of work to scan it in the backends.

Rather than create entirely new ModuleSummaryIndex structures (and all
the included unique_ptrs) for each backend index file, a map is created
to record all of the GUID and summary pointers needed for a particular
index file. The IndexBitcodeWriter walks this map instead of the full
index (hiding the details of managing the appropriate summary iteration
in a new iterator subclass). This is more efficient than walking the
entire combined index and filtering out just the needed summaries during
each backend bitcode index write.

Depends on D19481.

Reviewers: joker.eph

Subscribers: llvm-commits, joker.eph

Differential Revision: http://reviews.llvm.org/D19556

llvm-svn: 268627

1 files changed, 209 insertions, 71 deletions

diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
index 9fb5d66..6c3cecd 100644
--- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -263,6 +263,10 @@ class IndexBitcodeWriter : public BitcodeWriter {
   /// The combined index to write to bitcode.
   const ModuleSummaryIndex &Index;
 
+  /// When writing a subset of the index for distributed backends, client
+  /// provides a map of modules to the corresponding GUIDs/summaries to write.
+  std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex;
+
   /// Map that holds the correspondence between the GUID used in the combined
   /// index and a value id generated by this class to use in references.
   std::map<GlobalValue::GUID, unsigned> GUIDToValueIdMap;
@@ -272,17 +276,141 @@ class IndexBitcodeWriter : public BitcodeWriter {
 
 public:
   /// Constructs a IndexBitcodeWriter object for the given combined index,
-  /// writing to the provided \p Buffer.
+  /// writing to the provided \p Buffer. When writing a subset of the index
+  /// for a distributed backend, provide a \p ModuleToSummariesForIndex map.
   IndexBitcodeWriter(SmallVectorImpl<char> &Buffer,
-                     const ModuleSummaryIndex &Index)
-      : BitcodeWriter(Buffer), Index(Index) {
-    // Assign unique value ids to all functions in the index for use
+                     const ModuleSummaryIndex &Index,
+                     std::map<std::string, GVSummaryMapTy>
+                         *ModuleToSummariesForIndex = nullptr)
+      : BitcodeWriter(Buffer), Index(Index),
+        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
     // to the new global value id to use when writing those edges, which
     // are currently saved in the index in terms of GUID.
-    for (auto &II : Index)
-      GUIDToValueIdMap[II.first] = ++GlobalValueId;
-  }
+    for (const auto &I : *this)
+      GUIDToValueIdMap[I.first] = ++GlobalValueId;
+  }
+
+  /// The below iterator returns the GUID and associated summary.
+  typedef std::pair<GlobalValue::GUID, GlobalValueSummary *> GVInfo;
+
+  /// Iterator over the value GUID and summaries to be written to bitcode,
+  /// hides the details of whether they are being pulled from the entire
+  /// index or just those in a provided ModuleToSummariesForIndex map.
+  class iterator
+      : public llvm::iterator_facade_base<iterator, std::forward_iterator_tag,
+                                          GVInfo> {
+    /// Enables access to parent class.
+    const IndexBitcodeWriter &Writer;
+
+    // Iterators used when writing only those summaries in a provided
+    // ModuleToSummariesForIndex map:
+
+    /// Points to the last element in outer ModuleToSummariesForIndex map.
+    std::map<std::string, GVSummaryMapTy>::iterator ModuleSummariesBack;
+    /// Iterator on outer ModuleToSummariesForIndex map.
+    std::map<std::string, GVSummaryMapTy>::iterator ModuleSummariesIter;
+    /// Iterator on an inner global variable summary map.
+    GVSummaryMapTy::iterator ModuleGVSummariesIter;
+
+    // Iterators used when writing all summaries in the index:
+
+    /// Points to the last element in the Index outer GlobalValueMap.
+    const_gvsummary_iterator IndexSummariesBack;
+    /// Iterator on outer GlobalValueMap.
+    const_gvsummary_iterator IndexSummariesIter;
+    /// Iterator on an inner GlobalValueSummaryList.
+    GlobalValueSummaryList::const_iterator IndexGVSummariesIter;
+
+  public:
+    /// Construct iterator from parent \p Writer and indicate if we are
+    /// constructing the end iterator.
+    iterator(const IndexBitcodeWriter &Writer, bool IsAtEnd) : Writer(Writer) {
+      // Set up the appropriate set of iterators given whether we are writing
+      // the full index or just a subset.
+      // Can't setup the Back or inner iterators if the corresponding map
+      // is empty. This will be handled specially in operator== as well.
+      if (Writer.ModuleToSummariesForIndex &&
+          !Writer.ModuleToSummariesForIndex->empty()) {
+        ModuleSummariesBack =
+            std::prev(Writer.ModuleToSummariesForIndex->end());
+        ModuleSummariesIter = Writer.ModuleToSummariesForIndex->begin();
+        ModuleGVSummariesIter = !IsAtEnd ? ModuleSummariesIter->second.begin()
+                                         : ModuleSummariesBack->second.end();
+      } else if (!Writer.ModuleToSummariesForIndex &&
+                 Writer.Index.begin() != Writer.Index.end()) {
+        IndexSummariesBack = std::prev(Writer.Index.end());
+        IndexSummariesIter = Writer.Index.begin();
+        IndexGVSummariesIter = !IsAtEnd ? IndexSummariesIter->second.begin()
+                                        : IndexSummariesBack->second.end();
+      }
+    }
+
+    /// Increment the appropriate set of iterators.
+    iterator &operator++() {
+      // First the inner iterator is incremented, then if it is at the end
+      // and there are more outer iterations to go, the inner is reset to
+      // the start of the next inner list.
+      if (Writer.ModuleToSummariesForIndex) {
+        ++ModuleGVSummariesIter;
+        if (ModuleGVSummariesIter == ModuleSummariesIter->second.end() &&
+            ModuleSummariesIter != ModuleSummariesBack) {
+          ++ModuleSummariesIter;
+          ModuleGVSummariesIter = ModuleSummariesIter->second.begin();
+        }
+      } else {
+        ++IndexGVSummariesIter;
+        if (IndexGVSummariesIter == IndexSummariesIter->second.end() &&
+            IndexSummariesIter != IndexSummariesBack) {
+          ++IndexSummariesIter;
+          IndexGVSummariesIter = IndexSummariesIter->second.begin();
+        }
+      }
+      return *this;
+    }
+
+    /// Access the <GUID,GlobalValueSummary*> pair corresponding to the current
+    /// outer and inner iterator positions.
+    GVInfo operator*() {
+      if (Writer.ModuleToSummariesForIndex)
+        return std::make_pair(ModuleGVSummariesIter->first,
+                              ModuleGVSummariesIter->second);
+      return std::make_pair(IndexSummariesIter->first,
+                            IndexGVSummariesIter->get());
+    }
+
+    /// Checks if the iterators are equal, with special handling for empty
+    /// indexes.
+    bool operator==(const iterator &RHS) const {
+      if (Writer.ModuleToSummariesForIndex) {
+        // First ensure that both are writing the same subset.
+        if (Writer.ModuleToSummariesForIndex !=
+            RHS.Writer.ModuleToSummariesForIndex)
+          return false;
+        // Already determined above that maps are the same, so if one is
+        // empty, they both are.
+        if (Writer.ModuleToSummariesForIndex->empty())
+          return true;
+        return ModuleGVSummariesIter == RHS.ModuleGVSummariesIter;
+      }
+      // First ensure RHS also writing the full index, and that both are
+      // writing the same full index.
+      if (RHS.Writer.ModuleToSummariesForIndex ||
+          &Writer.Index != &RHS.Writer.Index)
+        return false;
+      // Already determined above that maps are the same, so if one is
+      // empty, they both are.
+      if (Writer.Index.begin() == Writer.Index.end())
+        return true;
+      return IndexGVSummariesIter == RHS.IndexGVSummariesIter;
+    }
+  };
+
+  /// Obtain the start iterator over the summaries to be written.
+  iterator begin() { return iterator(*this, /*IsAtEnd=*/false); }
+  /// Obtain the end iterator over the summaries to be written.
+  iterator end() { return iterator(*this, /*IsAtEnd=*/true); }
 
 private:
   /// Main entry point for writing a combined index to bitcode, invoked by
@@ -294,6 +422,14 @@ private:
   void writeCombinedValueSymbolTable();
   void writeCombinedGlobalValueSummary();
 
+  /// Indicates whether the provided \p ModulePath should be written into
+  /// the module string table, e.g. if full index written or if it is in
+  /// the provided subset.
+  bool doIncludeModule(StringRef ModulePath) {
+    return !ModuleToSummariesForIndex ||
+           ModuleToSummariesForIndex->count(ModulePath);
+  }
+
   bool hasValueId(GlobalValue::GUID ValGUID) {
     const auto &VMI = GUIDToValueIdMap.find(ValGUID);
     return VMI != GUIDToValueIdMap.end();
@@ -2963,6 +3099,8 @@ void IndexBitcodeWriter::writeModStrings() {
 
   SmallVector<unsigned, 64> Vals;
   for (const auto &MPSE : Index.modulePaths()) {
+    if (!doIncludeModule(MPSE.getKey()))
+      continue;
     StringEncoding Bits =
         getStringEncoding(MPSE.getKey().data(), MPSE.getKey().size());
     unsigned AbbrevToUse = Abbrev8Bit;
@@ -3218,78 +3356,75 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
     NameVals.clear();
   };
 
-  for (const auto &GSI : Index) {
-    for (auto &SI : GSI.second) {
-      GlobalValueSummary *S = SI.get();
-      assert(S);
-
-      assert(hasValueId(GSI.first));
-      unsigned ValueId = getValueId(GSI.first);
-      SummaryToValueIdMap[S] = ValueId;
-
-      if (auto *AS = dyn_cast<AliasSummary>(S)) {
-        // Will process aliases as a post-pass because the reader wants all
-        // global to be loaded first.
-        Aliases.push_back(AS);
-        continue;
-      }
+  for (const auto &I : *this) {
+    GlobalValueSummary *S = I.second;
+    assert(S);
 
-      if (auto *VS = dyn_cast<GlobalVarSummary>(S)) {
-        NameVals.push_back(ValueId);
-        NameVals.push_back(Index.getModuleId(VS->modulePath()));
-        NameVals.push_back(getEncodedGVSummaryFlags(VS->flags()));
-        for (auto &RI : VS->refs()) {
-          NameVals.push_back(getValueId(RI.getGUID()));
-        }
+    assert(hasValueId(I.first));
+    unsigned ValueId = getValueId(I.first);
+    SummaryToValueIdMap[S] = ValueId;
 
-        // Emit the finished record.
-        Stream.EmitRecord(bitc::FS_COMBINED_GLOBALVAR_INIT_REFS, NameVals,
-                          FSModRefsAbbrev);
-        NameVals.clear();
-        MaybeEmitOriginalName(*S);
-        continue;
-      }
+    if (auto *AS = dyn_cast<AliasSummary>(S)) {
+      // Will process aliases as a post-pass because the reader wants all
+      // global to be loaded first.
+      Aliases.push_back(AS);
+      continue;
+    }
 
-      auto *FS = cast<FunctionSummary>(S);
+    if (auto *VS = dyn_cast<GlobalVarSummary>(S)) {
       NameVals.push_back(ValueId);
-      NameVals.push_back(Index.getModuleId(FS->modulePath()));
-      NameVals.push_back(getEncodedGVSummaryFlags(FS->flags()));
-      NameVals.push_back(FS->instCount());
-      NameVals.push_back(FS->refs().size());
-
-      for (auto &RI : FS->refs()) {
+      NameVals.push_back(Index.getModuleId(VS->modulePath()));
+      NameVals.push_back(getEncodedGVSummaryFlags(VS->flags()));
+      for (auto &RI : VS->refs()) {
         NameVals.push_back(getValueId(RI.getGUID()));
       }
 
-      bool HasProfileData = false;
-      for (auto &EI : FS->calls()) {
-        HasProfileData |= EI.second.ProfileCount != 0;
-        if (HasProfileData)
-          break;
-      }
-
-      for (auto &EI : FS->calls()) {
-        // If this GUID doesn't have a value id, it doesn't have a function
-        // summary and we don't need to record any calls to it.
-        if (!hasValueId(EI.first.getGUID()))
-          continue;
-        NameVals.push_back(getValueId(EI.first.getGUID()));
-        assert(EI.second.CallsiteCount > 0 && "Expected at least one callsite");
-        NameVals.push_back(EI.second.CallsiteCount);
-        if (HasProfileData)
-          NameVals.push_back(EI.second.ProfileCount);
-      }
-
-      unsigned FSAbbrev =
-          (HasProfileData ? FSCallsProfileAbbrev : FSCallsAbbrev);
-      unsigned Code =
-          (HasProfileData ? bitc::FS_COMBINED_PROFILE : bitc::FS_COMBINED);
-
       // Emit the finished record.
-      Stream.EmitRecord(Code, NameVals, FSAbbrev);
+      Stream.EmitRecord(bitc::FS_COMBINED_GLOBALVAR_INIT_REFS, NameVals,
+                        FSModRefsAbbrev);
       NameVals.clear();
       MaybeEmitOriginalName(*S);
+      continue;
     }
+
+    auto *FS = cast<FunctionSummary>(S);
+    NameVals.push_back(ValueId);
+    NameVals.push_back(Index.getModuleId(FS->modulePath()));
+    NameVals.push_back(getEncodedGVSummaryFlags(FS->flags()));
+    NameVals.push_back(FS->instCount());
+    NameVals.push_back(FS->refs().size());
+
+    for (auto &RI : FS->refs()) {
+      NameVals.push_back(getValueId(RI.getGUID()));
+    }
+
+    bool HasProfileData = false;
+    for (auto &EI : FS->calls()) {
+      HasProfileData |= EI.second.ProfileCount != 0;
+      if (HasProfileData)
+        break;
+    }
+
+    for (auto &EI : FS->calls()) {
+      // If this GUID doesn't have a value id, it doesn't have a function
+      // summary and we don't need to record any calls to it.
+      if (!hasValueId(EI.first.getGUID()))
+        continue;
+      NameVals.push_back(getValueId(EI.first.getGUID()));
+      assert(EI.second.CallsiteCount > 0 && "Expected at least one callsite");
+      NameVals.push_back(EI.second.CallsiteCount);
+      if (HasProfileData)
+        NameVals.push_back(EI.second.ProfileCount);
+    }
+
+    unsigned FSAbbrev = (HasProfileData ? FSCallsProfileAbbrev : FSCallsAbbrev);
+    unsigned Code =
+        (HasProfileData ? bitc::FS_COMBINED_PROFILE : bitc::FS_COMBINED);
+
+    // Emit the finished record.
+    Stream.EmitRecord(Code, NameVals, FSAbbrev);
+    NameVals.clear();
+    MaybeEmitOriginalName(*S);
   }
 
   for (auto *AS : Aliases) {
@@ -3563,12 +3698,15 @@ void IndexBitcodeWriter::writeIndex() {
 
 // Write the specified module summary index to the given raw output stream,
 // where it will be written in a new bitcode block. This is used when
-// writing the combined index file for ThinLTO.
-void llvm::WriteIndexToFile(const ModuleSummaryIndex &Index, raw_ostream &Out) {
+// writing the combined index file for ThinLTO. When writing a subset of the
+// index for a distributed backend, provide a \p ModuleToSummariesForIndex map.
+void llvm::WriteIndexToFile(
+    const ModuleSummaryIndex &Index, raw_ostream &Out,
+    std::map<std::string, GVSummaryMapTy> *ModuleToSummariesForIndex) {
   SmallVector<char, 0> Buffer;
   Buffer.reserve(256 * 1024);
 
-  IndexBitcodeWriter IndexWriter(Buffer, Index);
+  IndexBitcodeWriter IndexWriter(Buffer, Index, ModuleToSummariesForIndex);
   IndexWriter.write();
 
   Out.write((char *)&Buffer.front(), Buffer.size());