From a30de6e7cb1458e271ead5e40e49dbe709275537 Mon Sep 17 00:00:00 2001
From: csmoe <csmoe@msn.com>
Date: Thu, 3 Nov 2022 22:34:24 +0800
Subject: [PATCH] record llvm cgu instruction stats

---
 Cargo.lock                                    |  2 ++
 compiler/rustc_codegen_llvm/Cargo.toml        |  2 ++
 compiler/rustc_codegen_llvm/src/back/write.rs | 21 +++++++++++++++++++
 compiler/rustc_codegen_llvm/src/llvm/ffi.rs   |  2 ++
 .../rustc_llvm/llvm-wrapper/LLVMWrapper.h     |  1 +
 .../rustc_llvm/llvm-wrapper/RustWrapper.cpp   | 13 ++++++++++++
 6 files changed, 41 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 3ffa9e590f402..d74b0a9aa5675 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3801,6 +3801,8 @@ dependencies = [
  "rustc_span",
  "rustc_symbol_mangling",
  "rustc_target",
+ "serde",
+ "serde_json",
  "smallvec",
  "tempfile",
  "tracing",
diff --git a/compiler/rustc_codegen_llvm/Cargo.toml b/compiler/rustc_codegen_llvm/Cargo.toml
index 773c0ebbe59db..a7ba2f8b69533 100644
--- a/compiler/rustc_codegen_llvm/Cargo.toml
+++ b/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"
diff --git a/compiler/rustc_codegen_llvm/src/back/write.rs b/compiler/rustc_codegen_llvm/src/back/write.rs
index 40f0594b40db3..a4ae1b01e869d 100644
--- a/compiler/rustc_codegen_llvm/src/back/write.rs
+++ b/compiler/rustc_codegen_llvm/src/back/write.rs
@@ -761,6 +761,7 @@ pub(crate) unsafe fn codegen(
             EmitObj::None => {}
         }
 
+        record_llvm_cgu_instructions_stats(&cgcx.prof, llmod);
         drop(handlers);
     }
 
@@ -974,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,
+        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);
+}
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
index 7aab666fc5e8c..1b3ce2e83a979 100644
--- a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
+++ b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -2410,6 +2410,8 @@ extern "C" {
     pub fn LLVMRustModuleBufferLen(p: &ModuleBuffer) -> usize;
     pub fn LLVMRustModuleBufferFree(p: &'static mut ModuleBuffer);
     pub fn LLVMRustModuleCost(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);
diff --git a/compiler/rustc_llvm/llvm-wrapper/LLVMWrapper.h b/compiler/rustc_llvm/llvm-wrapper/LLVMWrapper.h
index 9146a3739b2b1..0589062837866 100644
--- a/compiler/rustc_llvm/llvm-wrapper/LLVMWrapper.h
+++ b/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"
diff --git a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
index b1e6534944db3..e3493caaaf74e 100644
--- a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
+++ b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
@@ -1751,6 +1751,19 @@ LLVMRustModuleCost(LLVMModuleRef M) {
   return std::distance(std::begin(f), std::end(f));
 }
 
+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());
+  });
+}
+
 // Vector reductions:
 extern "C" LLVMValueRef
 LLVMRustBuildVectorReduceFAdd(LLVMBuilderRef B, LLVMValueRef Acc, LLVMValueRef Src) {