record llvm cgu instruction stats

Author: csmoe

Cargo.lock

@@ -3801,6 +3801,8 @@ dependencies = [
  "rustc_span",
  "rustc_symbol_mangling",
  "rustc_target",
+ "serde",
+ "serde_json",
  "smallvec",
  "tempfile",
  "tracing",
@@ -5186,9 +5188,9 @@ checksum = "5e9f0ab6ef7eb7353d9119c170a436d1bf248eea575ac42d19d12f4e34130831"
 
 [[package]]
 name = "snapbox"
-version = "0.4.1"
+version = "0.4.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "827c00e91b15e2674d8a5270bae91f898693cbf9561cbb58d8eaa31974597293"
+checksum = "09c20d8ee8713199cfd44148b70e47cb94f3e8dc538d727d31788f49e67b623e"
 dependencies = [
  "concolor",
  "content_inspector",

compiler/rustc_codegen_llvm/Cargo.toml

@@ -36,3 +36,5 @@ smallvec = { version = "1.8.1", features = ["union", "may_dangle"] }
 rustc_ast = { path = "../rustc_ast" }
 rustc_span = { path = "../rustc_span" }
 tempfile = "3.2.0"
+serde = { version = "1", features = [ "derive" ]}
+serde_json = "1"

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -20,6 +20,7 @@ use rustc_codegen_ssa::back::write::{
 };
 use rustc_codegen_ssa::traits::*;
 use rustc_codegen_ssa::{CompiledModule, ModuleCodegen};
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::profiling::SelfProfilerRef;
 use rustc_data_structures::small_c_str::SmallCStr;
 use rustc_errors::{FatalError, Handler, Level};
@@ -568,10 +569,8 @@ pub(crate) unsafe fn codegen(
         let llcx = &*module.module_llvm.llcx;
         let tm = &*module.module_llvm.tm;
         let module_name = module.name.clone();
-        let count = unsafe { llvm::LLVMRustModuleInstructionStats(&llmod) };
-        println!("llvm-module: {module_name} = {count}");
         let module_name = Some(&module_name[..]);
-        let handlers = DiagnosticHandlers::new(cgcx, diag_handler, llcx);
+        let _handlers = DiagnosticHandlers::new(cgcx, diag_handler, llcx);
 
         if cgcx.msvc_imps_needed {
             create_msvc_imps(cgcx, llcx, llmod);
@@ -763,7 +762,7 @@ pub(crate) unsafe fn codegen(
             EmitObj::None => {}
         }
 
-        drop(handlers);
+        record_llvm_cgu_instructions_stats(&cgcx.prof, llmod);
     }
 
     // `.dwo` files are only emitted if:
@@ -976,3 +975,23 @@ fn record_artifact_size(
         self_profiler_ref.artifact_size(artifact_kind, artifact_name.to_string_lossy(), file_size);
     }
 }
+
+fn record_llvm_cgu_instructions_stats(prof: &SelfProfilerRef, llmod: &llvm::Module) {
+    if !prof.enabled() {
+        return;
+    }
+
+    let raw_stats =
+        llvm::build_string(|s| unsafe { llvm::LLVMRustModuleInstructionStats(&llmod, s) })
+            .expect("cannot get module instruction stats");
+    #[derive(serde::Deserialize)]
+    struct InstructionsStats {
+        module: String,
+        _functions: FxHashMap<String, u64>,
+        total: u64,
+    }
+
+    let InstructionsStats { module, total, .. } =
+        serde_json::from_str(&raw_stats).expect("cannot parse llvm cgu instructions stats");
+    prof.artifact_size("cgu_instructions", module, total);
+}

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -2410,7 +2410,8 @@ extern "C" {
     pub fn LLVMRustModuleBufferLen(p: &ModuleBuffer) -> usize;
     pub fn LLVMRustModuleBufferFree(p: &'static mut ModuleBuffer);
     pub fn LLVMRustModuleCost(M: &Module) -> u64;
-    pub fn LLVMRustModuleInstructionStats(M: &Module) -> u64;
+    #[allow(improper_ctypes)]
+    pub fn LLVMRustModuleInstructionStats(M: &Module, Str: &RustString);
 
     pub fn LLVMRustThinLTOBufferCreate(M: &Module, is_thin: bool) -> &'static mut ThinLTOBuffer;
     pub fn LLVMRustThinLTOBufferFree(M: &'static mut ThinLTOBuffer);

compiler/rustc_llvm/llvm-wrapper/LLVMWrapper.h

@@ -14,6 +14,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/FormattedStream.h"
+#include "llvm/Support/JSON.h"
 #include "llvm/Support/Host.h"
 #include "llvm/Support/Memory.h"
 #include "llvm/Support/SourceMgr.h"

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -1751,18 +1751,26 @@ LLVMRustModuleCost(LLVMModuleRef M) {
   return std::distance(std::begin(f), std::end(f));
 }
 
-extern "C" uint64_t
-LLVMRustModuleInstructionStats(LLVMModuleRef M) {
-  auto f = unwrap(M)->functions();
-  raw_fd_ostream OS(2, false); // stderr.
-  for (auto &func : f)
-  {
-    auto name = func.getName();
-
-    auto count = func.getInstructionCount();
-    OS << name << "\t" << count << "\n";
-  }
-  return unwrap(M)->getInstructionCount();
+extern "C" void
+LLVMRustModuleInstructionStats(LLVMModuleRef M, RustStringRef Str)
+{
+  RawRustStringOstream OS(Str);
+  llvm::json::OStream JOS(OS);
+  auto Module = unwrap(M);
+
+  JOS.object([&]
+             {
+    JOS.attribute("module", Module->getName());
+    JOS.attribute("total", Module->getInstructionCount());
+    JOS.attributeArray("functions", [&] {
+        auto functions = Module->functions();
+        for (auto &f: functions) {
+            JOS.object([&] {
+                  JOS.attribute(f.getName(),f.getInstructionCount());
+            });
+        }
+
+    }); });
 }
 
 // Vector reductions: