Store item size estimate in MonoItemData.

This means we call `MonoItem::size_estimate` (which involves a query) less often: just once per mono item, and then once more per inline item placement. After that we can reuse the stored value as necessary. This means `CodegenUnit::compute_size_estimate` is cheaper.
crablang · Jul 16, 2023 · edd1f38 · edd1f38
1 parent b52f9eb
commit edd1f38
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 18 deletions.
diff --git a/compiler/rustc_middle/src/mir/mono.rs b/compiler/rustc_middle/src/mir/mono.rs
@@ -243,6 +243,7 @@ pub struct CodegenUnit<'tcx> {
 pub struct MonoItemData {
     pub linkage: Linkage,
     pub visibility: Visibility,
+    pub size_estimate: usize,
 }
 
 /// Specifies the linkage type for a `MonoItem`.
@@ -327,16 +328,16 @@ impl<'tcx> CodegenUnit<'tcx> {
         base_n::encode(hash, base_n::CASE_INSENSITIVE)
     }
 
-    pub fn compute_size_estimate(&mut self, tcx: TyCtxt<'tcx>) {
-        // Estimate the size of a codegen unit as (approximately) the number of MIR
-        // statements it corresponds to.
-        self.size_estimate = self.items.keys().map(|mi| mi.size_estimate(tcx)).sum();
+    pub fn compute_size_estimate(&mut self) {
+        // The size of a codegen unit as the sum of the sizes of the items
+        // within it.
+        self.size_estimate = self.items.values().map(|data| data.size_estimate).sum();
     }
 
-    #[inline]
     /// Should only be called if [`compute_size_estimate`] has previously been called.
     ///
     /// [`compute_size_estimate`]: Self::compute_size_estimate
+    #[inline]
     pub fn size_estimate(&self) -> usize {
         // Items are never zero-sized, so if we have items the estimate must be
         // non-zero, unless we forgot to call `compute_size_estimate` first.

diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs
@@ -222,11 +222,12 @@ where
     for mono_item in mono_items {
         // Handle only root items directly here. Inlined items are handled at
         // the bottom of the loop based on reachability.
+        let size_estimate = mono_item.size_estimate(cx.tcx);
         match mono_item.instantiation_mode(cx.tcx) {
             InstantiationMode::GloballyShared { .. } => {}
             InstantiationMode::LocalCopy => {
                 num_unique_inlined_items += 1;
-                unique_inlined_items_size += mono_item.size_estimate(cx.tcx);
+                unique_inlined_items_size += size_estimate;
                 continue;
             }
         }
@@ -258,7 +259,7 @@ where
             internalization_candidates.insert(mono_item);
         }
 
-        cgu.items_mut().insert(mono_item, MonoItemData { linkage, visibility });
+        cgu.items_mut().insert(mono_item, MonoItemData { linkage, visibility, size_estimate });
 
         // Get all inlined items that are reachable from `mono_item` without
         // going via another root item. This includes drop-glue, functions from
@@ -272,9 +273,11 @@ where
         // the `insert` will be a no-op.
         for inlined_item in reachable_inlined_items {
             // This is a CGU-private copy.
-            let linkage = Linkage::Internal;
-            let visibility = Visibility::Default;
-            cgu.items_mut().insert(inlined_item, MonoItemData { linkage, visibility });
+            cgu.items_mut().entry(inlined_item).or_insert_with(|| MonoItemData {
+                linkage: Linkage::Internal,
+                visibility: Visibility::Default,
+                size_estimate: inlined_item.size_estimate(cx.tcx),
+            });
         }
     }
 
@@ -289,7 +292,7 @@ where
     codegen_units.sort_by(|a, b| a.name().as_str().cmp(b.name().as_str()));
 
     for cgu in codegen_units.iter_mut() {
-        cgu.compute_size_estimate(cx.tcx);
+        cgu.compute_size_estimate();
     }
 
     return PlacedMonoItems {
@@ -352,7 +355,7 @@ fn merge_codegen_units<'tcx>(
             && codegen_units.iter().any(|cgu| cgu.size_estimate() < NON_INCR_MIN_CGU_SIZE))
     {
         // Sort small cgus to the back.
-        codegen_units.sort_by_cached_key(|cgu| cmp::Reverse(cgu.size_estimate()));
+        codegen_units.sort_by_key(|cgu| cmp::Reverse(cgu.size_estimate()));
 
         let mut smallest = codegen_units.pop().unwrap();
         let second_smallest = codegen_units.last_mut().unwrap();
@@ -361,7 +364,7 @@ fn merge_codegen_units<'tcx>(
         // may be duplicate inlined items, in which case the destination CGU is
         // unaffected. Recalculate size estimates afterwards.
         second_smallest.items_mut().extend(smallest.items_mut().drain());
-        second_smallest.compute_size_estimate(cx.tcx);
+        second_smallest.compute_size_estimate();
 
         // Record that `second_smallest` now contains all the stuff that was
         // in `smallest` before.
@@ -882,15 +885,15 @@ fn debug_dump<'a, 'tcx: 'a>(
             num_cgus += 1;
             all_cgu_sizes.push(cgu.size_estimate());
 
-            for (item, _) in cgu.items() {
+            for (item, data) in cgu.items() {
                 match item.instantiation_mode(tcx) {
                     InstantiationMode::GloballyShared { .. } => {
                         root_items += 1;
-                        root_size += item.size_estimate(tcx);
+                        root_size += data.size_estimate;
                     }
                     InstantiationMode::LocalCopy => {
                         placed_inlined_items += 1;
-                        placed_inlined_size += item.size_estimate(tcx);
+                        placed_inlined_size += data.size_estimate;
                     }
                 }
             }
@@ -932,7 +935,7 @@ fn debug_dump<'a, 'tcx: 'a>(
             let mean_size = size as f64 / num_items as f64;
 
             let mut placed_item_sizes: Vec<_> =
-                cgu.items().iter().map(|(item, _)| item.size_estimate(tcx)).collect();
+                cgu.items().values().map(|data| data.size_estimate).collect();
             placed_item_sizes.sort_unstable_by_key(|&n| cmp::Reverse(n));
             let sizes = list(&placed_item_sizes);
 
@@ -946,11 +949,11 @@ fn debug_dump<'a, 'tcx: 'a>(
                 let symbol_name = item.symbol_name(tcx).name;
                 let symbol_hash_start = symbol_name.rfind('h');
                 let symbol_hash = symbol_hash_start.map_or("<no hash>", |i| &symbol_name[i..]);
-                let size = item.size_estimate(tcx);
                 let kind = match item.instantiation_mode(tcx) {
                     InstantiationMode::GloballyShared { .. } => "root",
                     InstantiationMode::LocalCopy => "inlined",
                 };
+                let size = data.size_estimate;
                 let _ = with_no_trimmed_paths!(writeln!(
                     s,
                     "  - {item} [{linkage:?}] [{symbol_hash}] ({kind}, size: {size})"