Auto merge of #65281 - andjo403:partitioning, r=<try>

for a more even partitioning inline before merge

consider the size of the inlined items for a more even partitioning

for me this change take the compile time for script-servo-opt from 306s to 249s

edit: the times is for a 32 thread cpu.

cc #64913

Author: bors

Diff for: src/librustc/mir/mono.rs

@@ -47,20 +47,6 @@ pub enum MonoItem<'tcx> {
 }
 
 impl<'tcx> MonoItem<'tcx> {
-    pub fn size_estimate(&self, tcx: TyCtxt<'tcx>) -> usize {
-        match *self {
-            MonoItem::Fn(instance) => {
-                // Estimate the size of a function based on how many statements
-                // it contains.
-                tcx.instance_def_size_estimate(instance.def)
-            },
-            // Conservatively estimate the size of a static declaration
-            // or assembly to be 1.
-            MonoItem::Static(_) |
-            MonoItem::GlobalAsm(_) => 1,
-        }
-    }
-
     pub fn is_generic_fn(&self) -> bool {
         match *self {
             MonoItem::Fn(ref instance) => {
@@ -248,7 +234,7 @@ pub struct CodegenUnit<'tcx> {
     /// as well as the crate name and disambiguator.
     name: InternedString,
     items: FxHashMap<MonoItem<'tcx>, (Linkage, Visibility)>,
-    size_estimate: Option<usize>,
+    size_estimate: usize,
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, RustcEncodable, RustcDecodable)]
@@ -298,7 +284,7 @@ impl<'tcx> CodegenUnit<'tcx> {
         CodegenUnit {
             name: name,
             items: Default::default(),
-            size_estimate: None,
+            size_estimate: 0,
         }
     }
 
@@ -330,21 +316,17 @@ impl<'tcx> CodegenUnit<'tcx> {
         base_n::encode(hash, base_n::CASE_INSENSITIVE)
     }
 
-    pub fn estimate_size(&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 = Some(self.items.keys().map(|mi| mi.size_estimate(tcx)).sum());
-    }
-
     pub fn size_estimate(&self) -> usize {
-        // Should only be called if `estimate_size` has previously been called.
-        self.size_estimate.expect("estimate_size must be called before getting a size_estimate")
+        self.size_estimate
     }
 
-    pub fn modify_size_estimate(&mut self, delta: usize) {
-        assert!(self.size_estimate.is_some());
-        if let Some(size_estimate) = self.size_estimate {
-            self.size_estimate = Some(size_estimate + delta);
+    pub fn add_item(
+            &mut self,
+            item: MonoItem<'tcx>,
+            linkage: (Linkage, Visibility),
+            item_size: usize ) {
+        if self.items.insert(item, linkage).is_none() {
+            self.size_estimate += item_size;
         }
     }
 

Diff for: src/librustc_mir/monomorphize/partitioning.rs

@@ -136,25 +136,21 @@ where
 {
     let _prof_timer = tcx.prof.generic_activity("cgu_partitioning");
 
+    let mut estimator = SizeEstimator::new();
+
     // In the first step, we place all regular monomorphizations into their
     // respective 'home' codegen unit. Regular monomorphizations are all
     // functions and statics defined in the local crate.
-    let mut initial_partitioning = {
+    let initial_partitioning = {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_place_roots");
-        place_root_mono_items(tcx, mono_items)
+        place_root_mono_items(tcx, mono_items, &mut estimator)
     };
 
-    initial_partitioning.codegen_units.iter_mut().for_each(|cgu| cgu.estimate_size(tcx));
-
-    debug_dump(tcx, "INITIAL PARTITIONING:", initial_partitioning.codegen_units.iter());
-
-    // If the partitioning should produce a fixed count of codegen units, merge
-    // until that count is reached.
-    if let PartitioningStrategy::FixedUnitCount(count) = strategy {
-        let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_merge_cgus");
-        merge_codegen_units(tcx, &mut initial_partitioning, count);
-        debug_dump(tcx, "POST MERGING:", initial_partitioning.codegen_units.iter());
-    }
+    debug_dump(
+        tcx,
+        "INITIAL PARTITIONING:",
+        initial_partitioning.codegen_units.iter(),
+        &mut estimator);
 
     // In the next step, we use the inlining map to determine which additional
     // monomorphizations have to go into each codegen unit. These additional
@@ -163,12 +159,18 @@ where
     let mut post_inlining = {
         let _prof_timer =
             tcx.prof.generic_activity("cgu_partitioning_place_inline_items");
-        place_inlined_mono_items(initial_partitioning, inlining_map)
+        place_inlined_mono_items(tcx, initial_partitioning, inlining_map, &mut estimator)
     };
 
-    post_inlining.codegen_units.iter_mut().for_each(|cgu| cgu.estimate_size(tcx));
+    debug_dump(tcx, "POST INLINING:", post_inlining.codegen_units.iter(), &mut estimator);
 
-    debug_dump(tcx, "POST INLINING:", post_inlining.codegen_units.iter());
+    // If the partitioning should produce a fixed count of codegen units, merge
+    // until that count is reached.
+    if let PartitioningStrategy::FixedUnitCount(count) = strategy {
+        let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_merge_cgus");
+        merge_codegen_units(tcx, &mut post_inlining, count, &mut estimator);
+        debug_dump(tcx, "POST MERGING:", post_inlining.codegen_units.iter(), &mut estimator);
+    }
 
     // Next we try to make as many symbols "internal" as possible, so LLVM has
     // more freedom to optimize.
@@ -181,7 +183,6 @@ where
     // Finally, sort by codegen unit name, so that we get deterministic results.
     let PostInliningPartitioning {
         codegen_units: mut result,
-        mono_item_placements: _,
         internalization_candidates: _,
     } = post_inlining;
 
@@ -192,6 +193,27 @@ where
     result
 }
 
+struct SizeEstimator<'tcx> {
+    cache: FxHashMap<MonoItem<'tcx>, usize>,
+}
+impl<'tcx> SizeEstimator<'tcx> {
+    pub fn new() -> SizeEstimator<'tcx> {
+        SizeEstimator { cache: Default::default() }
+    }
+    pub fn size_estimate(&mut self, tcx: TyCtxt<'tcx>, mono_item: MonoItem<'tcx>) -> usize {
+        *self.cache.entry(mono_item).or_insert_with(|| match mono_item {
+            MonoItem::Fn(instance) => {
+                // Estimate the size of a function based on how many statements
+                // it contains.
+                tcx.instance_def_size_estimate(instance.def)
+            }
+            // Conservatively estimate the size of a static declaration
+            // or assembly to be 1.
+            MonoItem::Static(_) | MonoItem::GlobalAsm(_) => 1,
+        })
+    }
+}
+
 struct PreInliningPartitioning<'tcx> {
     codegen_units: Vec<CodegenUnit<'tcx>>,
     roots: FxHashSet<MonoItem<'tcx>>,
@@ -209,11 +231,14 @@ enum MonoItemPlacement {
 
 struct PostInliningPartitioning<'tcx> {
     codegen_units: Vec<CodegenUnit<'tcx>>,
-    mono_item_placements: FxHashMap<MonoItem<'tcx>, MonoItemPlacement>,
     internalization_candidates: FxHashSet<MonoItem<'tcx>>,
 }
 
-fn place_root_mono_items<'tcx, I>(tcx: TyCtxt<'tcx>, mono_items: I) -> PreInliningPartitioning<'tcx>
+fn place_root_mono_items<'tcx, I>(
+    tcx: TyCtxt<'tcx>,
+    mono_items: I,
+    estimator: &mut SizeEstimator<'tcx>,
+) -> PreInliningPartitioning<'tcx>
 where
     I: Iterator<Item = MonoItem<'tcx>>,
 {
@@ -264,8 +289,8 @@ where
         if visibility == Visibility::Hidden && can_be_internalized {
             internalization_candidates.insert(mono_item);
         }
-
-        codegen_unit.items_mut().insert(mono_item, (linkage, visibility));
+        let item_size = estimator.size_estimate(tcx, mono_item);
+        codegen_unit.add_item(mono_item, (linkage, visibility), item_size);
         roots.insert(mono_item);
     }
 
@@ -477,11 +502,12 @@ fn default_visibility(tcx: TyCtxt<'_>, id: DefId, is_generic: bool) -> Visibilit
 
 fn merge_codegen_units<'tcx>(
     tcx: TyCtxt<'tcx>,
-    initial_partitioning: &mut PreInliningPartitioning<'tcx>,
+    partitioning: &mut PostInliningPartitioning<'tcx>,
     target_cgu_count: usize,
+    estimator: &mut SizeEstimator<'tcx>,
 ) {
     assert!(target_cgu_count >= 1);
-    let codegen_units = &mut initial_partitioning.codegen_units;
+    let codegen_units = &mut partitioning.codegen_units;
 
     // Note that at this point in time the `codegen_units` here may not be in a
     // deterministic order (but we know they're deterministically the same set).
@@ -501,9 +527,9 @@ fn merge_codegen_units<'tcx>(
         let mut smallest = codegen_units.pop().unwrap();
         let second_smallest = codegen_units.last_mut().unwrap();
 
-        second_smallest.modify_size_estimate(smallest.size_estimate());
-        for (k, v) in smallest.items_mut().drain() {
-            second_smallest.items_mut().insert(k, v);
+        for (mono_item, v) in smallest.items_mut().drain() {
+            let item_size = estimator.size_estimate(tcx, mono_item);
+            second_smallest.add_item(mono_item, v, item_size);
         }
         debug!("CodegenUnit {} merged in to CodegenUnit {}",
                smallest.name(),
@@ -516,20 +542,19 @@ fn merge_codegen_units<'tcx>(
     }
 }
 
-fn place_inlined_mono_items<'tcx>(initial_partitioning: PreInliningPartitioning<'tcx>,
-                                  inlining_map: &InliningMap<'tcx>)
-                                  -> PostInliningPartitioning<'tcx> {
+fn place_inlined_mono_items<'tcx>(tcx: TyCtxt<'tcx>,
+                                  initial_partitioning: PreInliningPartitioning<'tcx>,
+                                  inlining_map: &InliningMap<'tcx>,
+                                  estimator: &mut SizeEstimator<'tcx>,
+) -> PostInliningPartitioning<'tcx> {
     let mut new_partitioning = Vec::new();
-    let mut mono_item_placements = FxHashMap::default();
 
     let PreInliningPartitioning {
         codegen_units: initial_cgus,
         roots,
         internalization_candidates,
     } = initial_partitioning;
 
-    let single_codegen_unit = initial_cgus.len() == 1;
-
     for old_codegen_unit in initial_cgus {
         // Collect all items that need to be available in this codegen unit.
         let mut reachable = FxHashSet::default();
@@ -543,40 +568,20 @@ fn place_inlined_mono_items<'tcx>(initial_partitioning: PreInliningPartitioning<
         for mono_item in reachable {
             if let Some(linkage) = old_codegen_unit.items().get(&mono_item) {
                 // This is a root, just copy it over.
-                new_codegen_unit.items_mut().insert(mono_item, *linkage);
+                let item_size = estimator.size_estimate(tcx, mono_item);
+                new_codegen_unit.add_item(mono_item, *linkage, item_size);
             } else {
                 if roots.contains(&mono_item) {
                     bug!("GloballyShared mono-item inlined into other CGU: \
                           {:?}", mono_item);
                 }
 
                 // This is a CGU-private copy.
-                new_codegen_unit.items_mut().insert(
+                let item_size = estimator.size_estimate(tcx, mono_item);
+                new_codegen_unit.add_item(
                     mono_item,
                     (Linkage::Internal, Visibility::Default),
-                );
-            }
-
-            if !single_codegen_unit {
-                // If there is more than one codegen unit, we need to keep track
-                // in which codegen units each monomorphization is placed.
-                match mono_item_placements.entry(mono_item) {
-                    Entry::Occupied(e) => {
-                        let placement = e.into_mut();
-                        debug_assert!(match *placement {
-                            MonoItemPlacement::SingleCgu { ref cgu_name } => {
-                                *cgu_name != *new_codegen_unit.name()
-                            }
-                            MonoItemPlacement::MultipleCgus => true,
-                        });
-                        *placement = MonoItemPlacement::MultipleCgus;
-                    }
-                    Entry::Vacant(e) => {
-                        e.insert(MonoItemPlacement::SingleCgu {
-                            cgu_name: new_codegen_unit.name().clone()
-                        });
-                    }
-                }
+                    item_size);
             }
         }
 
@@ -585,7 +590,6 @@ fn place_inlined_mono_items<'tcx>(initial_partitioning: PreInliningPartitioning<
 
     return PostInliningPartitioning {
         codegen_units: new_partitioning,
-        mono_item_placements,
         internalization_candidates,
     };
 
@@ -632,7 +636,24 @@ fn internalize_symbols<'tcx>(
         }
     });
 
-    let mono_item_placements = &partitioning.mono_item_placements;
+    // If there is more than one codegen unit, we need to keep track
+    // in which codegen units each monomorphization is placed.
+    let mut mono_item_placements = FxHashMap::default();
+
+    for cgu in &partitioning.codegen_units {
+        for mono_item in cgu.items().keys() {
+            match mono_item_placements.entry(*mono_item) {
+                Entry::Occupied(e) => {
+                    *e.into_mut() = MonoItemPlacement::MultipleCgus;
+                }
+                Entry::Vacant(e) => {
+                    e.insert(MonoItemPlacement::SingleCgu {
+                        cgu_name: cgu.name().clone()
+                    });
+                }
+            }
+        }
+    }
 
     // For each internalization candidates in each codegen unit, check if it is
     // accessed from outside its defining codegen unit.
@@ -781,7 +802,10 @@ fn numbered_codegen_unit_name(
     name_builder.build_cgu_name_no_mangle(LOCAL_CRATE, &["cgu"], Some(index))
 }
 
-fn debug_dump<'a, 'tcx, I>(tcx: TyCtxt<'tcx>, label: &str, cgus: I)
+fn debug_dump<'a, 'tcx, I>(tcx: TyCtxt<'tcx>,
+    label: &str,
+    cgus: I,
+    estimator: &mut SizeEstimator<'tcx>)
 where
     I: Iterator<Item = &'a CodegenUnit<'tcx>>,
     'tcx: 'a,
@@ -801,7 +825,7 @@ where
                        mono_item.to_string(tcx, true),
                        linkage,
                        symbol_hash,
-                       mono_item.size_estimate(tcx));
+                       estimator.size_estimate(tcx, *mono_item));
             }
 
             debug!("");