[ThinLTO] Fix printing of module paths for distributed backend indexes

Summary:
In the individual index files emitted for distributed ThinLTO backends,
the module path ids are not contiguous. Assign slots to module paths in
order to handle this better and also to get contiguous numbering in the
summary assembly.

Reviewers: davidxl, dexonsmith

Subscribers: mehdi_amini, inglorion, eraman, llvm-commits, steven_wu

Differential Revision: https://reviews.llvm.org/D48698

llvm-svn: 336148

llvm/lib/IR/AsmWriter.cpp

@@ -696,6 +696,10 @@ class SlotTracker {
   DenseMap<AttributeSet, unsigned> asMap;
   unsigned asNext = 0;
 
+  /// ModulePathMap - The slot map for Module paths used in the summary index.
+  StringMap<unsigned> ModulePathMap;
+  unsigned ModulePathNext = 0;
+
   /// GUIDMap - The slot map for GUIDs used in the summary index.
   DenseMap<GlobalValue::GUID, unsigned> GUIDMap;
   unsigned GUIDNext = 0;
@@ -729,6 +733,7 @@ class SlotTracker {
   int getGlobalSlot(const GlobalValue *V);
   int getMetadataSlot(const MDNode *N);
   int getAttributeGroupSlot(AttributeSet AS);
+  int getModulePathSlot(StringRef Path);
   int getGUIDSlot(GlobalValue::GUID GUID);
 
   /// If you'd like to deal with a function instead of just a module, use
@@ -782,6 +787,7 @@ class SlotTracker {
   /// Insert the specified AttributeSet into the slot table.
   void CreateAttributeSetSlot(AttributeSet AS);
 
+  inline void CreateModulePathSlot(StringRef Path);
   void CreateGUIDSlot(GlobalValue::GUID GUID);
 
   /// Add all of the module level global variables (and their initializers)
@@ -1005,8 +1011,16 @@ void SlotTracker::processIndex() {
   assert(TheIndex);
 
   // The first block of slots are just the module ids, which start at 0 and are
-  // assigned consecutively. Start numbering the GUIDs after the module ids.
-  GUIDNext = TheIndex->modulePaths().size();
+  // assigned consecutively. Since the StringMap iteration order isn't
+  // guaranteed, use a std::map to order by module ID before assigning slots.
+  std::map<uint64_t, StringRef> ModuleIdToPathMap;
+  for (auto &ModPath : TheIndex->modulePaths())
+    ModuleIdToPathMap[ModPath.second.first] = ModPath.first();
+  for (auto &ModPair : ModuleIdToPathMap)
+    CreateModulePathSlot(ModPair.second);
+
+  // Start numbering the GUIDs after the module ids.
+  GUIDNext = ModulePathNext;
 
   for (auto &GlobalList : *TheIndex)
     CreateGUIDSlot(GlobalList.first);
@@ -1100,6 +1114,15 @@ int SlotTracker::getAttributeGroupSlot(AttributeSet AS) {
   return AI == asMap.end() ? -1 : (int)AI->second;
 }
 
+int SlotTracker::getModulePathSlot(StringRef Path) {
+  // Check for uninitialized state and do lazy initialization.
+  initializeIndex();
+
+  // Find the Module path in the map
+  auto I = ModulePathMap.find(Path);
+  return I == ModulePathMap.end() ? -1 : (int)I->second;
+}
+
 int SlotTracker::getGUIDSlot(GlobalValue::GUID GUID) {
   // Check for uninitialized state and do lazy initialization.
   initializeIndex();
@@ -1169,6 +1192,11 @@ void SlotTracker::CreateAttributeSetSlot(AttributeSet AS) {
   asMap[AS] = DestSlot;
 }
 
+/// Create a new slot for the specified Module
+void SlotTracker::CreateModulePathSlot(StringRef Path) {
+  ModulePathMap[Path] = ModulePathNext++;
+}
+
 /// Create a new slot for the specified GUID
 void SlotTracker::CreateGUIDSlot(GlobalValue::GUID GUID) {
   GUIDMap[GUID] = GUIDNext++;
@@ -2574,23 +2602,19 @@ void AssemblyWriter::printModuleSummaryIndex() {
 
   Out << "\n";
 
-  // Print module path entries, using the module id as the slot number. To
-  // print in order, add paths to a vector indexed by module id.
+  // Print module path entries. To print in order, add paths to a vector
+  // indexed by module slot.
   std::vector<std::pair<std::string, ModuleHash>> moduleVec;
   std::string RegularLTOModuleName = "[Regular LTO]";
   moduleVec.resize(TheIndex->modulePaths().size());
-  for (auto &ModPath : TheIndex->modulePaths()) {
-    // A module id of -1 is a special entry for a regular LTO module created
-    // during the thin link.
-    if (ModPath.second.first == -1u)
-      moduleVec[TheIndex->modulePaths().size() - 1] =
-          std::make_pair(RegularLTOModuleName, ModPath.second.second);
-    else {
-      assert(ModPath.second.first < moduleVec.size());
-      moduleVec[ModPath.second.first] =
-          std::make_pair(ModPath.first(), ModPath.second.second);
-    }
-  }
+  for (auto &ModPath : TheIndex->modulePaths())
+    moduleVec[Machine.getModulePathSlot(ModPath.first())] = std::make_pair(
+        // A module id of -1 is a special entry for a regular LTO module created
+        // during the thin link.
+        ModPath.second.first == -1u ? RegularLTOModuleName
+                                    : (std::string)ModPath.first(),
+        ModPath.second.second);
+
   unsigned i = 0;
   for (auto &ModPair : moduleVec) {
     Out << "^" << i++ << " = module: (";
@@ -2937,7 +2961,7 @@ void AssemblyWriter::printSummary(const GlobalValueSummary &Summary) {
   GlobalValueSummary::GVFlags GVFlags = Summary.flags();
   GlobalValue::LinkageTypes LT = (GlobalValue::LinkageTypes)GVFlags.Linkage;
   Out << getSummaryKindName(Summary.getSummaryKind()) << ": ";
-  Out << "(module: ^" << TheIndex->getModuleId(Summary.modulePath())
+  Out << "(module: ^" << Machine.getModulePathSlot(Summary.modulePath())
       << ", flags: (";
   Out << "linkage: " << getLinkageName(LT);
   Out << ", notEligibleToImport: " << GVFlags.NotEligibleToImport;

llvm/test/tools/gold/X86/thinlto.ll

@@ -26,6 +26,8 @@
 ; RUN:    -shared %t.o %t2.o -o %t3
 ; RUN: llvm-bcanalyzer -dump %t.o.thinlto.bc | FileCheck %s --check-prefix=BACKEND1
 ; RUN: llvm-bcanalyzer -dump %t2.o.thinlto.bc | FileCheck %s --check-prefix=BACKEND2
+; RUN: llvm-dis %t.o.thinlto.bc -o - | FileCheck %s --check-prefix=DIS1
+; RUN: llvm-dis %t2.o.thinlto.bc -o - | FileCheck %s --check-prefix=DIS2
 ; RUN: not test -e %t3
 
 ; Ensure gold generates an index as well as a binary with save-temps in ThinLTO mode.
@@ -114,6 +116,14 @@
 ; BACKEND2-NEXT: <COMBINED
 ; BACKEND2-NEXT: </GLOBALVAL_SUMMARY_BLOCK
 
+; DIS1: ^0 = module: (path: "{{.*}}thinlto.ll.tmp.o", hash: (0, 0, 0, 0, 0))
+; DIS1: ^1 = module: (path: "{{.*}}thinlto.ll.tmp2.o", hash: (0, 0, 0, 0, 0))
+; DIS1: ^2 = gv: (guid: 13146401226427987378, summaries: (function: (module: ^1, flags: (linkage: external, notEligibleToImport: 0, live: 1, dsoLocal: 0), insts: 1)))
+; DIS1: ^3 = gv: (guid: 14740650423002898831, summaries: (function: (module: ^0, flags: (linkage: external, notEligibleToImport: 0, live: 1, dsoLocal: 0), insts: 2, calls: ((callee: ^2)))))
+
+; DIS2: ^0 = module: (path: "/usr/local/google/home/tejohnson/llvm/llvm_16_build/test/tools/gold/X86/Output/thinlto.ll.tmp2.o", hash: (0, 0, 0, 0, 0))
+; DIS2: ^1 = gv: (guid: 13146401226427987378, summaries: (function: (module: ^0, flags: (linkage: external, notEligibleToImport: 0, live: 1, dsoLocal: 0), insts: 1)))
+
 ; COMBINED: <MODULE_STRTAB_BLOCK
 ; COMBINED-NEXT: <ENTRY {{.*}} record string = '{{.*}}/test/tools/gold/X86/Output/thinlto.ll.tmp{{.*}}.o'
 ; COMBINED-NEXT: <ENTRY {{.*}} record string = '{{.*}}/test/tools/gold/X86/Output/thinlto.ll.tmp{{.*}}.o'