Add support for artifact size profiling

7 files changed, 57 insertions, 10 deletions

diff --git a/compiler/rustc_codegen_llvm/Cargo.toml b/compiler/rustc_codegen_llvm/Cargo.toml
index a6a553b31a3..d39c0bed5f8 100644
--- a/compiler/rustc_codegen_llvm/Cargo.toml
+++ b/compiler/rustc_codegen_llvm/Cargo.toml
@@ -11,7 +11,7 @@ doctest = false
 bitflags = "1.0"
 cstr = "0.2"
 libc = "0.2"
-measureme = "9.1.0"
+measureme = { git = "https://github.com/rylev/measureme" }
 snap = "1"
 tracing = "0.1"
 rustc_middle = { path = "../rustc_middle" }
diff --git a/compiler/rustc_data_structures/Cargo.toml b/compiler/rustc_data_structures/Cargo.toml
index 49962570129..b67329de2e8 100644
--- a/compiler/rustc_data_structures/Cargo.toml
+++ b/compiler/rustc_data_structures/Cargo.toml
@@ -23,7 +23,7 @@ rustc-hash = "1.1.0"
 smallvec = { version = "1.6.1", features = ["union", "may_dangle"] }
 rustc_index = { path = "../rustc_index", package = "rustc_index" }
 bitflags = "1.2.1"
-measureme = "9.1.0"
+measureme = { git = "https://github.com/rylev/measureme" }
 libc = "0.2"
 stacker = "0.1.14"
 tempfile = "3.2"
diff --git a/compiler/rustc_data_structures/src/profiling.rs b/compiler/rustc_data_structures/src/profiling.rs
index 0bbd0eda0c6..c21939209fc 100644
--- a/compiler/rustc_data_structures/src/profiling.rs
+++ b/compiler/rustc_data_structures/src/profiling.rs
@@ -110,12 +110,14 @@ bitflags::bitflags! {
         const FUNCTION_ARGS       = 1 << 6;
         const LLVM                = 1 << 7;
         const INCR_RESULT_HASHING = 1 << 8;
+        const ARTIFACT_SIZES = 1 << 9;
 
         const DEFAULT = Self::GENERIC_ACTIVITIES.bits |
                         Self::QUERY_PROVIDERS.bits |
                         Self::QUERY_BLOCKED.bits |
                         Self::INCR_CACHE_LOADS.bits |
-                        Self::INCR_RESULT_HASHING.bits;
+                        Self::INCR_RESULT_HASHING.bits |
+                        Self::ARTIFACT_SIZES.bits;
 
         const ARGS = Self::QUERY_KEYS.bits | Self::FUNCTION_ARGS.bits;
     }
@@ -136,6 +138,7 @@ const EVENT_FILTERS_BY_NAME: &[(&str, EventFilter)] = &[
     ("args", EventFilter::ARGS),
     ("llvm", EventFilter::LLVM),
     ("incr-result-hashing", EventFilter::INCR_RESULT_HASHING),
+    ("artifact-sizes", EventFilter::ARTIFACT_SIZES),
 ];
 
 /// Something that uniquely identifies a query invocation.
@@ -285,6 +288,33 @@ impl SelfProfilerRef {
         })
     }
 
+    /// Record the size of an artifact that the compiler produces
+    ///
+    /// `artifact_kind` is the class of artifact (e.g., query_cache, object_file, etc.)
+    /// `artifact_name` is an identifier to the specific artifact being stored (usually a filename)
+    #[inline(always)]
+    pub fn artifact_size<A>(&self, artifact_kind: &str, artifact_name: A, size: u64)
+    where
+        A: Borrow<str> + Into<String>,
+    {
+        drop(self.exec(EventFilter::ARTIFACT_SIZES, |profiler| {
+            let builder = EventIdBuilder::new(&profiler.profiler);
+            let event_label = profiler.get_or_alloc_cached_string(artifact_kind);
+            let event_arg = profiler.get_or_alloc_cached_string(artifact_name);
+            let event_id = builder.from_label_and_arg(event_label, event_arg);
+            let thread_id = get_thread_id();
+
+            profiler.profiler.record_integer_event(
+                profiler.artifact_size_event_kind,
+                event_id,
+                thread_id,
+                size,
+            );
+
+            TimingGuard::none()
+        }))
+    }
+
     #[inline(always)]
     pub fn generic_activity_with_args(
         &self,
@@ -372,7 +402,7 @@ impl SelfProfilerRef {
     ) {
         drop(self.exec(event_filter, |profiler| {
             let event_id = StringId::new_virtual(query_invocation_id.0);
-            let thread_id = std::thread::current().id().as_u64().get() as u32;
+            let thread_id = get_thread_id();
 
             profiler.profiler.record_instant_event(
                 event_kind(profiler),
@@ -425,6 +455,7 @@ pub struct SelfProfiler {
     incremental_result_hashing_event_kind: StringId,
     query_blocked_event_kind: StringId,
     query_cache_hit_event_kind: StringId,
+    artifact_size_event_kind: StringId,
 }
 
 impl SelfProfiler {
@@ -447,6 +478,7 @@ impl SelfProfiler {
             profiler.alloc_string("IncrementalResultHashing");
         let query_blocked_event_kind = profiler.alloc_string("QueryBlocked");
         let query_cache_hit_event_kind = profiler.alloc_string("QueryCacheHit");
+        let artifact_size_event_kind = profiler.alloc_string("ArtifactSize");
 
         let mut event_filter_mask = EventFilter::empty();
 
@@ -491,6 +523,7 @@ impl SelfProfiler {
             incremental_result_hashing_event_kind,
             query_blocked_event_kind,
             query_cache_hit_event_kind,
+            artifact_size_event_kind,
         })
     }
 
@@ -561,7 +594,7 @@ impl<'a> TimingGuard<'a> {
         event_kind: StringId,
         event_id: EventId,
     ) -> TimingGuard<'a> {
-        let thread_id = std::thread::current().id().as_u64().get() as u32;
+        let thread_id = get_thread_id();
         let raw_profiler = &profiler.profiler;
         let timing_guard =
             raw_profiler.start_recording_interval_event(event_kind, event_id, thread_id);
@@ -655,6 +688,10 @@ pub fn duration_to_secs_str(dur: std::time::Duration) -> String {
     format!("{:.3}", dur.as_secs_f64())
 }
 
+fn get_thread_id() -> u32 {
+    std::thread::current().id().as_u64().get() as u32
+}
+
 // Memory reporting
 cfg_if! {
     if #[cfg(windows)] {
diff --git a/compiler/rustc_incremental/src/persist/file_format.rs b/compiler/rustc_incremental/src/persist/file_format.rs
index 572a4fc6971..392c5bdc15a 100644
--- a/compiler/rustc_incremental/src/persist/file_format.rs
+++ b/compiler/rustc_incremental/src/persist/file_format.rs
@@ -95,6 +95,12 @@ where
         return;
     }
 
+    sess.prof.artifact_size(
+        &name.replace(' ', "_"),
+        path_buf.file_name().unwrap().to_string_lossy(),
+        encoder.position() as u64,
+    );
+
     debug!("save: data written to disk successfully");
 }
 
diff --git a/compiler/rustc_query_impl/Cargo.toml b/compiler/rustc_query_impl/Cargo.toml
index 89df3d4674b..c4da929e298 100644
--- a/compiler/rustc_query_impl/Cargo.toml
+++ b/compiler/rustc_query_impl/Cargo.toml
@@ -7,7 +7,7 @@ edition = "2021"
 doctest = false
 
 [dependencies]
-measureme = "9.0.0"
+measureme = { git = "https://github.com/rylev/measureme" }
 rustc-rayon-core = "0.3.1"
 tracing = "0.1"
 rustc_ast = { path = "../rustc_ast" }
diff --git a/compiler/rustc_query_system/src/dep_graph/serialized.rs b/compiler/rustc_query_system/src/dep_graph/serialized.rs
index f5f67fcd0a0..47197a1e492 100644
--- a/compiler/rustc_query_system/src/dep_graph/serialized.rs
+++ b/compiler/rustc_query_system/src/dep_graph/serialized.rs
@@ -222,7 +222,7 @@ impl<K: DepKind> EncoderState<K> {
         index
     }
 
-    fn finish(self) -> FileEncodeResult {
+    fn finish(self, profiler: &SelfProfilerRef) -> FileEncodeResult {
         let Self { mut encoder, total_node_count, total_edge_count, result, stats: _ } = self;
         let () = result?;
 
@@ -235,7 +235,11 @@ impl<K: DepKind> EncoderState<K> {
         IntEncodedWithFixedSize(edge_count).encode(&mut encoder)?;
         debug!("position: {:?}", encoder.position());
         // Drop the encoder so that nothing is written after the counts.
-        encoder.flush()
+        let result = encoder.flush();
+        // FIXME(rylev): we hardcode the dep graph file name so we don't need a dependency on
+        // rustc_incremental just for that.
+        profiler.artifact_size("dep_graph", "dep-graph.bin", encoder.position() as u64);
+        result
     }
 }
 
@@ -332,6 +336,6 @@ impl<K: DepKind + Encodable<FileEncoder>> GraphEncoder<K> {
 
     pub fn finish(self, profiler: &SelfProfilerRef) -> FileEncodeResult {
         let _prof_timer = profiler.generic_activity("incr_comp_encode_dep_graph");
-        self.status.into_inner().finish()
+        self.status.into_inner().finish(profiler)
     }
 }
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index 8ecb7a031ad..b7f4396f9d7 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -1279,7 +1279,7 @@ options! {
         "specify the events recorded by the self profiler;
         for example: `-Z self-profile-events=default,query-keys`
         all options: none, all, default, generic-activity, query-provider, query-cache-hit
-                     query-blocked, incr-cache-load, incr-result-hashing, query-keys, function-args, args, llvm"),
+                     query-blocked, incr-cache-load, incr-result-hashing, query-keys, function-args, args, llvm, artifact-sizes"),
     share_generics: Option<bool> = (None, parse_opt_bool, [TRACKED],
         "make the current crate share its generic instantiations"),
     show_span: Option<String> = (None, parse_opt_string, [TRACKED],