[Stats] Collect IR/LLVM always-on stats correctly when multithreading.

lib/FrontendTool/FrontendTool.cpp

@@ -67,7 +67,6 @@
 #include "llvm/ADT/Statistic.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
-#include "llvm/IR/ValueSymbolTable.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Option/Option.h"
 #include "llvm/Option/OptTable.h"
@@ -463,25 +462,6 @@ static void countStatsPostSema(UnifiedStatsReporter &Stats,
   }
 }
 
-static void countStatsPostIRGen(UnifiedStatsReporter &Stats,
-                                const llvm::Module& Module) {
-  auto &C = Stats.getFrontendCounters();
-  // FIXME: calculate these in constant time if possible.
-  C.NumIRGlobals = Module.getGlobalList().size();
-  C.NumIRFunctions = Module.getFunctionList().size();
-  C.NumIRAliases = Module.getAliasList().size();
-  C.NumIRIFuncs = Module.getIFuncList().size();
-  C.NumIRNamedMetaData = Module.getNamedMDList().size();
-  C.NumIRValueSymbols = Module.getValueSymbolTable().size();
-  C.NumIRComdatSymbols = Module.getComdatSymbolTable().size();
-  for (auto const &Func : Module) {
-    for (auto const &BB : Func) {
-      C.NumIRBasicBlocks++;
-      C.NumIRInsts += BB.size();
-    }
-  }
-}
-
 static void countStatsPostSILGen(UnifiedStatsReporter &Stats,
                                  const SILModule& Module) {
   auto &C = Stats.getFrontendCounters();
@@ -1027,10 +1007,6 @@ static bool performCompile(CompilerInstance &Instance,
     return HadError;
   }
 
-  if (Stats) {
-    countStatsPostIRGen(*Stats, *IRModule);
-  }
-
   bool allSymbols = false;
   switch (opts.ValidateTBDAgainstIR) {
   case FrontendOptions::TBDValidationMode::None:

lib/IRGen/IRGen.cpp

@@ -44,6 +44,7 @@
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/ValueSymbolTable.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Linker/Linker.h"
 #include "llvm/MC/SubtargetFeature.h"
@@ -355,6 +356,25 @@ static bool needsRecompile(StringRef OutputFilename, ArrayRef<uint8_t> HashData,
   return true;
 }
 
+static void countStatsPostIRGen(UnifiedStatsReporter &Stats,
+                                const llvm::Module& Module) {
+  auto &C = Stats.getFrontendCounters();
+  // FIXME: calculate these in constant time if possible.
+  C.NumIRGlobals += Module.getGlobalList().size();
+  C.NumIRFunctions += Module.getFunctionList().size();
+  C.NumIRAliases += Module.getAliasList().size();
+  C.NumIRIFuncs += Module.getIFuncList().size();
+  C.NumIRNamedMetaData += Module.getNamedMDList().size();
+  C.NumIRValueSymbols += Module.getValueSymbolTable().size();
+  C.NumIRComdatSymbols += Module.getComdatSymbolTable().size();
+  for (auto const &Func : Module) {
+    for (auto const &BB : Func) {
+      C.NumIRBasicBlocks++;
+      C.NumIRInsts += BB.size();
+    }
+  }
+}
+
 /// Run the LLVM passes. In multi-threaded compilation this will be done for
 /// multiple LLVM modules in parallel.
 bool swift::performLLVM(IRGenOptions &Opts, DiagnosticEngine *Diags,
@@ -466,6 +486,14 @@ bool swift::performLLVM(IRGenOptions &Opts, DiagnosticEngine *Diags,
   }
   }
 
+  if (Stats) {
+    if (DiagMutex)
+      DiagMutex->lock();
+    countStatsPostIRGen(*Stats, *Module);
+    if (DiagMutex)
+      DiagMutex->unlock();
+  }
+
   EmitPasses.run(*Module);
 
   if (Stats && RawOS.hasValue()) {
@@ -791,7 +819,7 @@ static std::unique_ptr<llvm::Module> performIRGeneration(IRGenOptions &Opts,
     if (performLLVM(Opts, &IGM.Context.Diags, nullptr, IGM.ModuleHash,
                     IGM.getModule(), IGM.TargetMachine.get(),
                     IGM.Context.LangOpts.EffectiveLanguageVersion,
-                    IGM.OutputFilename))
+                    IGM.OutputFilename, IGM.Context.Stats))
       return nullptr;
   }
 
@@ -811,7 +839,7 @@ static void ThreadEntryPoint(IRGenerator *irgen,
     performLLVM(irgen->Opts, &IGM->Context.Diags, DiagMutex, IGM->ModuleHash,
                 IGM->getModule(), IGM->TargetMachine.get(),
                 IGM->Context.LangOpts.EffectiveLanguageVersion,
-                IGM->OutputFilename);
+                IGM->OutputFilename, IGM->Context.Stats);
     if (IGM->Context.Diags.hadAnyError())
       return;
   }

test/Misc/stats_dir.swift

@@ -2,6 +2,10 @@
 // RUN: %target-swift-frontend -c -o %t/out.o -stats-output-dir %t %s
 // RUN: %utils/process-stats-dir.py --set-csv-baseline %t/frontend.csv %t
 // RUN: %FileCheck -input-file %t/frontend.csv %s
+// RUN: rm -rf %t && mkdir -p %t
+// RUN: %target-swift-frontend -c -wmo -num-threads 4 -o %t/out.o -stats-output-dir %t %s
+// RUN: %utils/process-stats-dir.py --set-csv-baseline %t/frontend.csv %t
+// RUN: %FileCheck -input-file %t/frontend.csv %s
 // RUN: echo '9000000000	"LLVM.NumLLVMBytesOutput"	1' >>%t/frontend.csv
 // RUN: not %utils/process-stats-dir.py --compare-to-csv-baseline %t/frontend.csv %t
 
@@ -11,7 +15,8 @@
 // RUN: %FileCheck -input-file %t/driver.csv %s
 // RUN: %utils/process-stats-dir.py --compare-to-csv-baseline %t/driver.csv %t
 
-// CHECK: LLVM.NumLLVMBytesOutput
+// CHECK: {{"IRModule.NumIRFunctions"	[1-9][0-9]*$}}
+// CHECK: {{"LLVM.NumLLVMBytesOutput"	[1-9][0-9]*$}}
 
 public func foo() {
     print("hello")