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

[tejohnson]

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.

[dexonsmith]

What was the numbering before?

[davidxl]

lib/IR/AsmWriter.cpp

790

make it an inline function.

1019

Can these two loops be fused ?

1119

probably better to rename the function as initializeIndexIfNeeded or something, but can be done separately.

[tejohnson]

lib/IR/AsmWriter.cpp

1019

No, there are 2 loops for the reason noted in the comments: Since the StringMap iteration order isn't

// guaranteed, use a std::map to order by module ID before assigning slots.

1119

Will do as a follow-on (I named this similar to the existing "initialize()" which does lazy initialization of slots needed for printing the Module.

[tejohnson]

Address comments

[davidxl]

lgtm

lib/IR/AsmWriter.cpp

790

CreateModulePathSlot is a tiny wrapper which should be marked as inline. I saw a different function is marked instead.

[tejohnson]

lib/IR/AsmWriter.cpp

790

Woops, will fix.

[tejohnson]

Mark correct function as inline

lib/IR/AsmWriter.cpp

private:
   /// mdnMap - Map for MDNodes.
   DenseMap<const MDNode*, unsigned> mdnMap;
   unsigned mdnNext = 0;
 
   /// asMap - The slot map for attribute sets.
   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;
 
 public:
   /// Construct from a module.
   ///
   /// If \c ShouldInitializeAllMetadata, initializes all metadata in all
public:
   SlotTracker &operator=(const SlotTracker &) = delete;
 
   /// Return the slot number of the specified value in it's type
   /// plane.  If something is not in the SlotTracker, return -1.
   int getLocalSlot(const Value *V);
   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
   /// this method to get its data into the SlotTracker.
   void incorporateFunction(const Function *F) {
     TheFunction = F;
     FunctionProcessed = false;
   }
public:
 
   /// These functions do the actual initialization.
   inline void initialize();
   void initializeIndex();
 
   // Implementation Details
 private:
   /// CreateModuleSlot - Insert the specified GlobalValue* into the slot table.
-  void CreateModuleSlot(const GlobalValue *V);
+  inline void CreateModuleSlot(const GlobalValue *V);
 
   /// CreateMetadataSlot - Insert the specified MDNode* into the slot table.
   void CreateMetadataSlot(const MDNode *N);
 
   /// CreateFunctionSlot - Insert the specified Value* into the slot table.
   void CreateFunctionSlot(const Value *V);
 
   /// Insert the specified AttributeSet into the slot table.
   void CreateAttributeSetSlot(AttributeSet AS);
 
+  void CreateModulePathSlot(StringRef Path);
   void CreateGUIDSlot(GlobalValue::GUID GUID);
 
   /// Add all of the module level global variables (and their initializers)
   /// and function declarations, but not the contents of those functions.
   void processModule();
   void processIndex();
 
   /// Add all of the functions arguments, basic blocks, and instructions.

 }
 
 // Iterate through all the GUID in the index and create slots for them.
 void SlotTracker::processIndex() {
   ST_DEBUG("begin processIndex!\n");
   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);
 
   for (auto &TId : TheIndex->typeIds())
     CreateGUIDSlot(GlobalValue::getGUID(TId.first));
 
   ST_DEBUG("end processIndex!\n");
int SlotTracker::getAttributeGroupSlot(AttributeSet AS) {
   // Check for uninitialized state and do lazy initialization.
   initialize();
 
   // Find the AttributeSet in the module map.
   as_iterator AI = asMap.find(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();
 
   // Find the GUID in the map
   guid_iterator I = GUIDMap.find(GUID);
   return I == GUIDMap.end() ? -1 : (int)I->second;
 }
void SlotTracker::CreateAttributeSetSlot(AttributeSet AS) {
   as_iterator I = asMap.find(AS);
   if (I != asMap.end())
     return;
 
   unsigned DestSlot = asNext++;
   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++;
 }
 
 //===----------------------------------------------------------------------===//
 // AsmWriter Implementation
 //===----------------------------------------------------------------------===//

 }
 
 void AssemblyWriter::printModuleSummaryIndex() {
   assert(TheIndex);
   Machine.initializeIndex();
 
   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()) {
+  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.
-    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);
-    }
-  }
+        ModPath.second.first == -1u ? RegularLTOModuleName
+                                    : (std::string)ModPath.first(),
+        ModPath.second.second);
+
   unsigned i = 0;
   for (auto &ModPair : moduleVec) {
     Out << "^" << i++ << " = module: (";
     Out << "path: \"" << ModPair.first << "\"";
     Out << ", hash: (";
     FieldSeparator FS;
     for (auto Hash : ModPair.second)
       Out << FS << Hash;
void AssemblyWriter::printConstVCalls(
   }
   Out << ")";
 }
 
 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;
   Out << ", live: " << GVFlags.Live;
   Out << ", dsoLocal: " << GVFlags.DSOLocal;
   Out << ")";
 
   if (Summary.getSummaryKind() == GlobalValueSummary::AliasKind)

test/tools/gold/X86/thinlto.ll

 ; Ensure gold generates an index and not a binary if requested.
 ; RUN: %gold -plugin %llvmshlibdir/LLVMgold%shlibext \
 ; RUN:    -m elf_x86_64 \
 ; RUN:    --plugin-opt=thinlto \
 ; RUN:    --plugin-opt=thinlto-index-only \
 ; 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.
 ; First force single-threaded mode
 ; RUN: %gold -plugin %llvmshlibdir/LLVMgold%shlibext \
 ; RUN:    -m elf_x86_64 \
 ; RUN:    --plugin-opt=save-temps \
 ; RUN:    --plugin-opt=thinlto \

 ; BACKEND2-NEXT: </MODULE_STRTAB_BLOCK
 ; BACKEND2-NEXT: <GLOBALVAL_SUMMARY_BLOCK
 ; BACKEND2-NEXT: <VERSION
 ; BACKEND2-NEXT: <FLAGS
 ; BACKEND2-NEXT: <VALUE_GUID op0=1 op1=-5300342847281564238
 ; 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'
 ; COMBINED-NEXT: </MODULE_STRTAB_BLOCK
 ; COMBINED-NEXT: <GLOBALVAL_SUMMARY_BLOCK
 ; COMBINED-NEXT: <VERSION
 ; COMBINED-NEXT: <FLAGS
 ; COMBINED-NEXT: <VALUE_GUID op0={{1|2}} op1={{-3706093650706652785|-5300342847281564238}}