[ThinLTO] Compute the basic block count across modules.

Summary:
Count the per-module number of basic blocks when the module summary is computed
and sum them up during Thin LTO indexing.

This is used to estimate the working set size under the partial sample PGO.

This is split off of D79831.

Reviewers: davidxl, espindola

Subscribers: emaste, inglorion, hiraditya, MaskRay, steven_wu, dexonsmith, arphaman, llvm-commits

Tags: #llvm

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

include/llvm/Bitcode/LLVMBitCodes.h

@@ -290,6 +290,8 @@ enum GlobalValueSummarySymtabCodes {
   //                                        numrefs, numrefs x valueid,
   //                                        n x (valueid, offset)]
   FS_PERMODULE_VTABLE_GLOBALVAR_INIT_REFS = 23,
+  // The total number of basic blocks in the module.
+  FS_BLOCK_COUNT = 24,
 };
 
 enum MetadataCodes {

include/llvm/IR/ModuleSummaryIndex.h

@@ -1007,6 +1007,10 @@ private:
   StringSaver Saver;
   BumpPtrAllocator Alloc;
 
+  // The total number of basic blocks in the module in the per-module summary or
+  // the total number of basic blocks in the LTO unit in the combined index.
+  uint64_t BlockCount;
+
   // YAML I/O support.
   friend yaml::MappingTraits<ModuleSummaryIndex>;
 
@@ -1019,8 +1023,8 @@ private:
 public:
   // See HaveGVs variable comment.
   ModuleSummaryIndex(bool HaveGVs, bool EnableSplitLTOUnit = false)
-      : HaveGVs(HaveGVs), EnableSplitLTOUnit(EnableSplitLTOUnit), Saver(Alloc) {
+      : HaveGVs(HaveGVs), EnableSplitLTOUnit(EnableSplitLTOUnit), Saver(Alloc),
-  }
+        BlockCount(0) {}
 
   // Current version for the module summary in bitcode files.
   // The BitcodeSummaryVersion should be bumped whenever we introduce changes
@@ -1039,6 +1043,10 @@ public:
   uint64_t getFlags() const;
   void setFlags(uint64_t Flags);
 
+  uint64_t getBlockCount() const { return BlockCount; }
+  void addBlockCount(uint64_t C) { BlockCount += C; }
+  void setBlockCount(uint64_t C) { BlockCount = C; }
+
   gvsummary_iterator begin() { return GlobalValueMap.begin(); }
   const_gvsummary_iterator begin() const { return GlobalValueMap.begin(); }
   gvsummary_iterator end() { return GlobalValueMap.end(); }

lib/Analysis/ModuleSummaryAnalysis.cpp

@@ -390,6 +390,7 @@ static void computeFunctionSummary(ModuleSummaryIndex &Index, const Module &M,
               .updateHotness(getHotness(Candidate.Count, PSI));
       }
     }
+  Index.addBlockCount(F.size());
 
   std::vector<ValueInfo> Refs;
   if (IsThinLTO) {

lib/AsmParser/LLLexer.cpp

@@ -741,6 +741,7 @@ lltok::Kind LLLexer::LexIdentifier() {
   KEYWORD(name);
   KEYWORD(summaries);
   KEYWORD(flags);
+  KEYWORD(blockcount);
   KEYWORD(linkage);
   KEYWORD(notEligibleToImport);
   KEYWORD(live);

lib/AsmParser/LLParser.cpp

@@ -853,6 +853,9 @@ bool LLParser::ParseSummaryEntry() {
   case lltok::kw_flags:
     result = ParseSummaryIndexFlags();
     break;
+  case lltok::kw_blockcount:
+    result = ParseBlockCount();
+    break;
   default:
     result = Error(Lex.getLoc(), "unexpected summary kind");
     break;
@@ -8111,6 +8114,21 @@ bool LLParser::ParseSummaryIndexFlags() {
   return false;
 }
 
+/// ParseBlockCount
+///   ::= 'blockcount' ':' UInt64
+bool LLParser::ParseBlockCount() {
+  assert(Lex.getKind() == lltok::kw_blockcount);
+  Lex.Lex();
+
+  if (ParseToken(lltok::colon, "expected ':' here"))
+    return true;
+  uint64_t BlockCount;
+  if (ParseUInt64(BlockCount))
+    return true;
+  Index->setBlockCount(BlockCount);
+  return false;
+}
+
 /// ParseGVEntry
 ///   ::= 'gv' ':' '(' ('name' ':' STRINGCONSTANT | 'guid' ':' UInt64)
 ///         [',' 'summaries' ':' Summary[',' Summary]* ]? ')'

lib/AsmParser/LLParser.h

@@ -341,6 +341,7 @@ namespace llvm {
     bool ParseModuleReference(StringRef &ModulePath);
     bool ParseGVReference(ValueInfo &VI, unsigned &GVId);
     bool ParseSummaryIndexFlags();
+    bool ParseBlockCount();
     bool ParseGVEntry(unsigned ID);
     bool ParseFunctionSummary(std::string Name, GlobalValue::GUID, unsigned ID);
     bool ParseVariableSummary(std::string Name, GlobalValue::GUID, unsigned ID);

lib/AsmParser/LLToken.h

@@ -374,6 +374,7 @@ enum Kind {
   kw_name,
   kw_summaries,
   kw_flags,
+  kw_blockcount,
   kw_linkage,
   kw_notEligibleToImport,
   kw_live,

lib/Bitcode/Reader/BitcodeAnalyzer.cpp

@@ -305,6 +305,7 @@ static Optional<const char *> GetCodeName(unsigned CodeID, unsigned BlockID,
       STRINGIFY_CODE(FS, CFI_FUNCTION_DECLS)
       STRINGIFY_CODE(FS, TYPE_ID)
       STRINGIFY_CODE(FS, TYPE_ID_METADATA)
+      STRINGIFY_CODE(FS, BLOCK_COUNT)
     }
   case bitc::METADATA_ATTACHMENT_ID:
     switch (CodeID) {

lib/Bitcode/Reader/BitcodeReader.cpp

@@ -6239,6 +6239,9 @@ Error ModuleSummaryIndexBitcodeReader::parseEntireSummary(unsigned ID) {
     case bitc::FS_TYPE_ID_METADATA:
       parseTypeIdCompatibleVtableSummaryRecord(Record);
       break;
+
+    case bitc::FS_BLOCK_COUNT:
+      TheIndex.addBlockCount(Record[0]);
     }
   }
   llvm_unreachable("Exit infinite loop");

lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -3906,6 +3906,9 @@ void ModuleBitcodeWriterBase::writePerModuleGlobalValueSummary() {
     NameVals.clear();
   }
 
+  Stream.EmitRecord(bitc::FS_BLOCK_COUNT,
+                    ArrayRef<uint64_t>{Index->getBlockCount()});
+
   Stream.ExitBlock();
 }
 
@@ -4189,6 +4192,9 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
     }
   }
 
+  Stream.EmitRecord(bitc::FS_BLOCK_COUNT,
+                    ArrayRef<uint64_t>{Index.getBlockCount()});
+
   Stream.ExitBlock();
 }
 

lib/IR/AsmWriter.cpp

@@ -2822,8 +2822,13 @@ void AssemblyWriter::printModuleSummaryIndex() {
   }
 
   // Don't emit flags when it's not really needed (value is zero by default).
-  if (TheIndex->getFlags())
+  if (TheIndex->getFlags()) {
     Out << "^" << NumSlots << " = flags: " << TheIndex->getFlags() << "\n";
+    ++NumSlots;
+  }
+
+  Out << "^" << NumSlots << " = blockcount: " << TheIndex->getBlockCount()
+      << "\n";
 }
 
 static const char *

test/Assembler/thinlto-vtable-summary.ll

@@ -36,3 +36,4 @@ declare i32 @_ZN1C1fEi(%struct.C*, i32)
 ^6 = typeidCompatibleVTable: (name: "_ZTS1A", summary: ((offset: 16, ^2), (offset: 16, ^4))) ; guid = 7004155349499253778
 ^7 = typeidCompatibleVTable: (name: "_ZTS1B", summary: ((offset: 16, ^2))) ; guid = 6203814149063363976
 ^8 = typeidCompatibleVTable: (name: "_ZTS1C", summary: ((offset: 16, ^4))) ; guid = 1884921850105019584
+^9 = blockcount: 0

test/Bitcode/thinlto-alias.ll

@@ -22,6 +22,7 @@
 ; See if the call to func is registered.
 ; The value id 1 matches the second FUNCTION record above.
 ; CHECK-NEXT:    <PERMODULE {{.*}} op7=1/>
+; CHECK-NEXT:    <BLOCK_COUNT op0=1/>
 ; CHECK-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 ; CHECK: <STRTAB_BLOCK
@@ -37,6 +38,7 @@
 ; COMBINED-NEXT:    <COMBINED {{.*}} op9=[[ALIASID]]/>
 ; COMBINED-NEXT:    <COMBINED {{.*}}
 ; COMBINED-NEXT:    <COMBINED_ALIAS  {{.*}} op3=[[ALIASEEID]]
+; COMBINED-NEXT:    <BLOCK_COUNT op0=2/>
 ; COMBINED-NEXT:  </GLOBALVAL_SUMMARY_BLOCK
 
 ; ModuleID = 'thinlto-function-summary-callgraph.ll'
@@ -55,9 +57,11 @@ declare void @analias(...)
 ; DIS: ^0 = module: (path: "{{.*}}", hash: (0, 0, 0, 0, 0))
 ; DIS: ^1 = gv: (name: "analias", summaries: (alias: (module: ^0, flags: (linkage: external, notEligibleToImport: 0, live: 0, dsoLocal: 0, canAutoHide: 0), aliasee: ^2))) ; guid = 12695095382722328222
 ; DIS: ^2 = gv: (name: "aliasee", summaries: (function: (module: ^0, flags: (linkage: external, notEligibleToImport: 0, live: 0, dsoLocal: 0, canAutoHide: 0), insts: 1))) ; guid = 17407585008595848568
+; DIS: ^3 = blockcount: 1
 
 ; COMBINED-DIS: ^0 = module: (path: "{{.*}}thinlto-alias.ll.tmp.o", hash: (0, 0, 0, 0, 0))
 ; COMBINED-DIS: ^1 = module: (path: "{{.*}}thinlto-alias.ll.tmp2.o", hash: (0, 0, 0, 0, 0))
 ; COMBINED-DIS: ^2 = gv: (guid: 12695095382722328222, summaries: (alias: (module: ^1, flags: (linkage: external, notEligibleToImport: 0, live: 0, dsoLocal: 0, canAutoHide: 0), aliasee: ^4)))
 ; COMBINED-DIS: ^3 = gv: (guid: 15822663052811949562, summaries: (function: (module: ^0, flags: (linkage: external, notEligibleToImport: 0, live: 0, dsoLocal: 0, canAutoHide: 0), insts: 2, calls: ((callee: ^2)))))
 ; COMBINED-DIS: ^4 = gv: (guid: 17407585008595848568, summaries: (function: (module: ^1, flags: (linkage: external, notEligibleToImport: 0, live: 0, dsoLocal: 0, canAutoHide: 0), insts: 1)))
+; COMBINED-DIS: ^5 = blockcount: 2

test/Bitcode/thinlto-alias2.ll

@@ -8,6 +8,7 @@
 ; CHECK-NEXT:    <PERMODULE {{.*}} op4=0 op5=0 op6=0 op7=[[ALIASID:[0-9]+]]/>
 ; CHECK-NEXT:    <PERMODULE {{.*}} op0=[[ALIASEEID:[0-9]+]]
 ; CHECK-NEXT:    <ALIAS {{.*}} op0=[[ALIASID]] {{.*}} op2=[[ALIASEEID]]/>
+; CHECK-NEXT:    <BLOCK_COUNT op0=2/>
 ; CHECK-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 ; ModuleID = 'thinlto-alias2.ll'

test/Bitcode/thinlto-function-summary-callgraph-cast.ll

@@ -12,6 +12,7 @@
 ; CHECK-NEXT:    <PERMODULE {{.*}} op4=3 {{.*}} op9={{[0-9]+}}/>
 ; CHECK-NEXT:    <PERMODULE {{.*}} op0=[[ALIASEEID:[0-9]+]]
 ; CHECK-NEXT:    <ALIAS {{.*}} op0=[[ALIASID]] {{.*}} op2=[[ALIASEEID]]/>
+; CHECK-NEXT:    <BLOCK_COUNT op0=3/>
 ; CHECK-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 ; ModuleID = 'thinlto-function-summary-callgraph-cast.ll'

test/Bitcode/thinlto-function-summary-callgraph-pgo.ll

@@ -19,6 +19,7 @@
 ; CHECK-NEXT:    <FLAGS
 ; See if the call to func is registered, using the expected hotness type.
 ; CHECK-NEXT:    <PERMODULE_PROFILE {{.*}} op7=1 op8=2/>
+; CHECK-NEXT:    <BLOCK_COUNT op0=1/>
 ; CHECK-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 ; CHECK: <STRTAB_BLOCK
 ; CHECK-NEXT: blob data = 'mainfunc{{.*}}'
@@ -32,6 +33,7 @@
 ; See if the call to func is registered, using the expected hotness type.
 ; op6=2 which is hotnessType::None.
 ; COMBINED-NEXT:    <COMBINED_PROFILE {{.*}} op9=[[FUNCID]] op10=2/>
+; COMBINED-NEXT:    <BLOCK_COUNT op0=2/>
 ; COMBINED-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 ; ModuleID = 'thinlto-function-summary-callgraph.ll'

test/Bitcode/thinlto-function-summary-callgraph-profile-summary.ll

@@ -50,6 +50,7 @@
 ; CHECK-NEXT:    <VALUE_GUID op0=25 op1=123/>
 ; op4=hot1 op6=cold op8=hot2 op10=hot4 op12=none1 op14=hot3 op16=none2 op18=none3 op20=123
 ; CHECK-NEXT:    <PERMODULE_PROFILE {{.*}} op7=1 op8=3 op9=5 op10=1 op11=2 op12=3 op13=4 op14=1 op15=6 op16=2 op17=3 op18=3 op19=7 op20=2 op21=8 op22=2 op23=25 op24=4/>
+; CHECK-NEXT:    <BLOCK_COUNT op0=6/>
 ; CHECK-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 ; CHECK: <STRTAB_BLOCK
@@ -73,8 +74,9 @@
 ; COMBINED-NEXT:    <COMBINED abbrevid=
 ; COMBINED-NEXT:    <COMBINED abbrevid=
 ; COMBINED-NEXT:    <COMBINED_PROFILE {{.*}} op9=[[HOT1:.*]] op10=3 op11=[[COLD:.*]] op12=1 op13=[[HOT2:.*]] op14=3 op15=[[NONE1:.*]] op16=2 op17=[[HOT3:.*]] op18=3 op19=[[NONE2:.*]] op20=2 op21=[[NONE3:.*]] op22=2/>
-; COMBINED_NEXT:    <COMBINED abbrevid=
+; COMBINED-NEXT:    <COMBINED abbrevid=
-; COMBINED_NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
+; COMBINED-NEXT:    <BLOCK_COUNT op0=13/>
+; COMBINED-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 
 ; ModuleID = 'thinlto-function-summary-callgraph.ll'

test/Bitcode/thinlto-function-summary-callgraph-relbf.ll

@@ -15,6 +15,7 @@
 ; CHECK-NEXT:    <FLAGS
 ; See if the call to func is registered.
 ; CHECK-NEXT:    <PERMODULE_RELBF {{.*}} op4=1 {{.*}} op9=256
+; CHECK-NEXT:    <BLOCK_COUNT op0=1/>
 ; CHECK-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 ; CHECK: <STRTAB_BLOCK
 ; CHECK-NEXT: blob data = 'undefinedglobmainfunc{{.*}}'

test/Bitcode/thinlto-function-summary-callgraph-sample-profile-summary.ll

@@ -33,6 +33,7 @@
 ; CHECK-NEXT:    <VALUE_GUID op0=26 op1=123/>
 ; op4=none1 op6=hot1 op8=cold1 op10=none2 op12=hot2 op14=cold2 op16=none3 op18=hot3 op20=cold3 op22=123
 ; CHECK-NEXT:    <PERMODULE_PROFILE {{.*}} op7=7 op8=0 op9=1 op10=3 op11=4 op12=1 op13=8 op14=0 op15=2 op16=3 op17=5 op18=1 op19=9 op20=0 op21=3 op22=3 op23=6 op24=1 op25=26 op26=4/>
+; CHECK-NEXT:    <BLOCK_COUNT op0=4/>
 ; CHECK-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 ; CHECK: <STRTAB_BLOCK
@@ -60,8 +61,9 @@
 ; COMBINED-NEXT:    <COMBINED abbrevid=
 ; COMBINED-NEXT:    <COMBINED abbrevid=
 ; COMBINED-NEXT:    <COMBINED_PROFILE {{.*}} op9=[[NONE1:.*]] op10=0 op11=[[HOT1:.*]] op12=3 op13=[[COLD1:.*]] op14=1 op15=[[NONE2:.*]] op16=0 op17=[[HOT2:.*]] op18=3 op19=[[COLD2:.*]] op20=1 op21=[[NONE3:.*]] op22=0 op23=[[HOT3:.*]] op24=3 op25=[[COLD3:.*]] op26=1/>
-; COMBINED_NEXT:    <COMBINED abbrevid=
+; COMBINED-NEXT:    <COMBINED abbrevid=
-; COMBINED_NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
+; COMBINED-NEXT:    <BLOCK_COUNT op0=13/>
+; COMBINED-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 
 ; ModuleID = 'thinlto-function-summary-callgraph.ll'

test/Bitcode/thinlto-function-summary-callgraph.ll

@@ -20,6 +20,7 @@
 ; CHECK-NEXT:    <FLAGS
 ; See if the call to func is registered
 ; CHECK-NEXT:    <PERMODULE {{.*}} op4=1
+; CHECK-NEXT:    <BLOCK_COUNT op0=1/>
 ; CHECK-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 ; CHECK: <STRTAB_BLOCK
 ; CHECK-NEXT: blob data = 'undefinedglobmainfunc{{.*}}'
@@ -35,6 +36,7 @@
 ; COMBINED-NEXT:    <COMBINED
 ; See if the call to func is registered.
 ; COMBINED-NEXT:    <COMBINED {{.*}} op9=[[FUNCID]]/>
+; COMBINED-NEXT:    <BLOCK_COUNT op0=2/>
 ; COMBINED-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 ; ModuleID = 'thinlto-function-summary-callgraph.ll'

test/Bitcode/thinlto-function-summary-originalnames.ll

@@ -15,6 +15,7 @@
 ; COMBINED-DAG:    <COMBINED_ORIGINAL_NAME op0=-2012135647395072713/>
 ; COMBINED-DAG:    <COMBINED_ALIAS
 ; COMBINED-DAG:    <COMBINED_ORIGINAL_NAME op0=-4170563161550796836/>
+; COMBINED-NEXT:   <BLOCK_COUNT op0=1/>
 ; COMBINED-NEXT:  </GLOBALVAL_SUMMARY_BLOCK>
 
 source_filename = "/path/to/source.c"

test/Bitcode/thinlto-function-summary.ll

@@ -25,6 +25,7 @@
 ; BC-NEXT: <PERMODULE {{.*}} op0=3 op1=7
 ; BC-NEXT: <PERMODULE {{.*}} op0=4 op1=0 op2=4 op3=0
 ; BC-NEXT: <ALIAS {{.*}} op0=6 op1=0 op2=3
+; BC-NEXT: <BLOCK_COUNT op0=5/>
 ; BC-NEXT: </GLOBALVAL_SUMMARY_BLOCK
 ; BC: <STRTAB_BLOCK
 ; BC-NEXT: blob data = 'hfoobaranon.{{................................}}.0variadicllvm.va_startf{{.*}}'

test/ThinLTO/X86/distributed_indexes.ll

@@ -23,6 +23,7 @@
 ; BACKEND1-DAG: <COMBINED {{.*}} op1=0
 ; BACKEND1-DAG: <COMBINED {{.*}} op1=1
 ; BACKEND1-DAG: <COMBINED_ALIAS {{.*}} op1=1
+; BACKEND1-NEXT: <BLOCK_COUNT op0=3/>
 ; BACKEND1-NEXT: </GLOBALVAL_SUMMARY_BLOCK
 
 ; The backend index for Input/distributed_indexes.ll contains summaries from
@@ -38,6 +39,7 @@
 ; BACKEND2-NEXT: <COMBINED
 ; BACKEND2-NEXT: <COMBINED
 ; BACKEND2-NEXT: <COMBINED_ALIAS
+; BACKEND2-NEXT: <BLOCK_COUNT op0=3/>
 ; BACKEND2-NEXT: </GLOBALVAL_SUMMARY_BLOCK
 
 ; Make sure that when the alias is imported as a copy of the aliasee, but the

test/tools/llvm-lto/thinlto.ll

@@ -16,6 +16,7 @@
 ; COMBINED-NEXT: <VALUE_GUID op0={{1|2}} op1={{-3706093650706652785|-5300342847281564238}}
 ; COMBINED-NEXT: <COMBINED
 ; COMBINED-NEXT: <COMBINED
+; COMBINED-NEXT: <BLOCK_COUNT op0=2/>
 ; COMBINED-NEXT: </GLOBALVAL_SUMMARY_BLOCK
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"