Rollup merge of #105481 - lqd:mono-stats, r=wesleywiser

Start improving monomorphization items stats

As described in [this zulip discussion](https://rust-lang.zulipchat.com/#narrow/stream/247081-t-compiler.2Fperformance/topic/Compile-time.20case-study.3A.20AWS.20crates/near/314560832), some stats about monomorphization collection would be interesting to have, in a different form than `-Zprint-mono-items`: to have some visibility into the cost of the mono items, we'd like to know how many are instantiated and what is their estimated size.

That can be a proxy to analyze sources of slow compile times, although in the future, we'd also like to add more realistic stats from the actual backend's lowering.

This PR adds a new `-Z dump-mono-stats` flag which will output some stats in a `{crate_name}.mono-items.md` file (the flag optionally takes an output directory parameter, for easier use within a workspace than printing to stdout).

For example,

```rust
fn compute<T>(collection: Vec<T>) -> usize {
    collection.len() + 19 - 0 * 9 - 18 - 1 * 1 // random code to increase the function's size
}

fn main() {
    dbg!(compute(vec![0u8, 1, 2]));
    dbg!(compute(vec![0u64, 1, 2]));
    dbg!(compute(vec!["0", "1", "2", "3"]));
}
```

will output a file with this markdown table (abridged for readability), for a debug build:

| Item | Instantiation count | Estimated Cost Per Instantiation | Total Estimated Cost |
| --- | ---: | ---: | ---: |
| alloc::alloc::box_free | 3 | 122 | 366 |
| std::alloc::Global::alloc_impl | 1 | 284 | 284 |
| alloc::raw_vec::RawVec::<T, A>::current_memory | 3 | 82 | 246 |
| std::ptr::align_offset | 1 | 222 | 222 |
| std::slice::hack::into_vec | 3 | 67 | 201 |
| <std::vec::Vec<T, A> as std::ops::Drop>::drop | 3 | 66 | 198 |
| std::ptr::mut_ptr::<impl *mut T>::is_null | 4 | 47 | 188 |
| main | 1 | 163 | 163 |
| std::ptr::NonNull::<T>::new_unchecked | 4 | 37 | 148 |
...

<details>
<summary>Click for full output</summary>

| Item | Instantiation count | Estimated Cost Per Instantiation | Total Estimated Cost |
| --- | ---: | ---: | ---: |
| alloc::alloc::box_free | 3 | 122 | 366 |
| std::alloc::Global::alloc_impl | 1 | 284 | 284 |
| alloc::raw_vec::RawVec::<T, A>::current_memory | 3 | 82 | 246 |
| std::ptr::align_offset | 1 | 222 | 222 |
| std::slice::hack::into_vec | 3 | 67 | 201 |
| <std::vec::Vec<T, A> as std::ops::Drop>::drop | 3 | 66 | 198 |
| std::ptr::mut_ptr::<impl *mut T>::is_null | 4 | 47 | 188 |
| main | 1 | 163 | 163 |
| std::ptr::NonNull::<T>::new_unchecked | 4 | 37 | 148 |
| std::boxed::Box::<T, A>::into_unique | 3 | 48 | 144 |
| std::boxed::Box::<T, A>::leak | 3 | 39 | 117 |
| std::alloc::Layout::array::inner | 1 | 107 | 107 |
| std::ptr::align_offset::mod_inv | 1 | 103 | 103 |
| std::boxed::Box::<T, A>::into_raw_with_allocator | 3 | 31 | 93 |
| std::fmt::Arguments::<'a>::new_v1 | 1 | 80 | 80 |
| <alloc::raw_vec::RawVec<T, A> as std::ops::Drop>::drop | 3 | 26 | 78 |
| alloc::raw_vec::RawVec::<T, A>::from_raw_parts_in | 3 | 26 | 78 |
| alloc::alloc::exchange_malloc | 1 | 75 | 75 |
| std::ptr::const_ptr::<impl *const T>::is_null | 1 | 75 | 75 |
| std::ptr::const_ptr::<impl *const T>::is_aligned_to | 1 | 64 | 64 |
| compute | 3 | 20 | 60 |
| std::ptr::const_ptr::<impl *const T>::align_offset | 1 | 55 | 55 |
| std::ptr::read | 1 | 52 | 52 |
| <std::alloc::Global as std::alloc::Allocator>::deallocate | 1 | 50 | 50 |
| std::ptr::mut_ptr::<impl *mut T>::guaranteed_eq | 1 | 48 | 48 |
| std::fmt::ArgumentV1::<'a>::new_display | 2 | 22 | 44 |
| std::ptr::Alignment::new_unchecked | 1 | 42 | 42 |
| core::fmt::num::<impl std::fmt::Debug for usize>::fmt | 1 | 40 | 40 |
| std::result::Result::<T, E>::unwrap_unchecked | 1 | 37 | 37 |
| std::cmp::Ord::min | 1 | 32 | 32 |
| std::cmp::impls::<impl std::cmp::Ord for usize>::cmp | 1 | 31 | 31 |
| std::intrinsics::is_aligned_and_not_null | 1 | 27 | 27 |
| std::rt::lang_start | 1 | 27 | 27 |
| std::ptr::NonNull::<T>::new | 1 | 24 | 24 |
| std::fmt::ArgumentV1::<'a>::new_debug | 1 | 22 | 22 |
| std::fmt::Arguments::<'a>::new_v1_formatted | 1 | 19 | 19 |
| std::rt::lang_start::{closure#0} | 1 | 17 | 17 |
| std::sys_common::backtrace::__rust_begin_short_backtrace | 1 | 16 | 16 |
| std::slice::<impl [T]>::into_vec | 3 | 5 | 15 |
| <std::ptr::NonNull<T> as std::convert::From<std::ptr::Unique<T>>>::from | 1 | 14 | 14 |
| <&T as std::fmt::Debug>::fmt | 1 | 12 | 12 |
| <&T as std::fmt::Display>::fmt | 1 | 12 | 12 |
| std::vec::Vec::<T, A>::len | 3 | 2 | 6 |
| <T as std::convert::Into<U>>::into | 1 | 5 | 5 |
| <T as std::convert::From<T>>::from | 1 | 2 | 2 |
| <() as std::process::Termination>::report | 1 | 2 | 2 |
| std::hint::unreachable_unchecked | 1 | 2 | 2 |
| core::fmt::UnsafeArg::new | 1 | 1 | 1 |

</details>

Since we discussed it together, r? `@wesleywiser.`

compiler/rustc_error_messages/locales/en-US/monomorphize.ftl

@@ -21,3 +21,6 @@ monomorphize_large_assignments =
     moving {$size} bytes
     .label = value moved from here
     .note = The current maximum size is {$limit}, but it can be customized with the move_size_limit attribute: `#![move_size_limit = "..."]`
+
+monomorphize_couldnt_dump_mono_stats =
+    unexpected error occurred while dumping monomorphization stats: {$error}

compiler/rustc_middle/src/mir/mono.rs

@@ -200,6 +200,15 @@ impl<'tcx> MonoItem<'tcx> {
             MonoItem::GlobalAsm(..) => LOCAL_CRATE,
         }
     }
+
+    /// Returns the item's `DefId`
+    pub fn def_id(&self) -> DefId {
+        match *self {
+            MonoItem::Fn(Instance { def, .. }) => def.def_id(),
+            MonoItem::Static(def_id) => def_id,
+            MonoItem::GlobalAsm(item_id) => item_id.owner_id.to_def_id(),
+        }
+    }
 }
 
 impl<'tcx> fmt::Display for MonoItem<'tcx> {

compiler/rustc_monomorphize/src/errors.rs

@@ -77,3 +77,9 @@ pub struct SymbolAlreadyDefined {
     pub span: Option<Span>,
     pub symbol: String,
 }
+
+#[derive(Diagnostic)]
+#[diag(monomorphize_couldnt_dump_mono_stats)]
+pub struct CouldntDumpMonoStats {
+    pub error: String,
+}

compiler/rustc_monomorphize/src/partitioning/mod.rs

@@ -95,6 +95,11 @@
 mod default;
 mod merging;
 
+use std::cmp;
+use std::fs::{self, File};
+use std::io::{BufWriter, Write};
+use std::path::{Path, PathBuf};
+
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::sync;
 use rustc_hir::def_id::DefIdSet;
@@ -104,11 +109,12 @@ use rustc_middle::mir::mono::{CodegenUnit, Linkage};
 use rustc_middle::ty::print::with_no_trimmed_paths;
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::TyCtxt;
+use rustc_session::config::SwitchWithOptPath;
 use rustc_span::symbol::Symbol;
 
 use crate::collector::InliningMap;
 use crate::collector::{self, MonoItemCollectionMode};
-use crate::errors::{SymbolAlreadyDefined, UnknownPartitionStrategy};
+use crate::errors::{CouldntDumpMonoStats, SymbolAlreadyDefined, UnknownPartitionStrategy};
 
 pub struct PartitioningCx<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
@@ -411,6 +417,15 @@ fn collect_and_partition_mono_items<'tcx>(
         })
         .collect();
 
+    // Output monomorphization stats per def_id
+    if let SwitchWithOptPath::Enabled(ref path) = tcx.sess.opts.unstable_opts.dump_mono_stats {
+        if let Err(err) =
+            dump_mono_items_stats(tcx, &codegen_units, path, tcx.sess.opts.crate_name.as_deref())
+        {
+            tcx.sess.emit_fatal(CouldntDumpMonoStats { error: err.to_string() });
+        }
+    }
+
     if tcx.sess.opts.unstable_opts.print_mono_items.is_some() {
         let mut item_to_cgus: FxHashMap<_, Vec<_>> = Default::default();
 
@@ -465,6 +480,67 @@ fn collect_and_partition_mono_items<'tcx>(
     (tcx.arena.alloc(mono_items), codegen_units)
 }
 
+/// Outputs stats about instantation counts and estimated size, per `MonoItem`'s
+/// def, to a file in the given output directory.
+fn dump_mono_items_stats<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    codegen_units: &[CodegenUnit<'tcx>],
+    output_directory: &Option<PathBuf>,
+    crate_name: Option<&str>,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let output_directory = if let Some(ref directory) = output_directory {
+        fs::create_dir_all(directory)?;
+        directory
+    } else {
+        Path::new(".")
+    };
+
+    let filename = format!("{}.mono_items.md", crate_name.unwrap_or("unknown-crate"));
+    let output_path = output_directory.join(&filename);
+    let file = File::create(output_path)?;
+    let mut file = BufWriter::new(file);
+
+    // Gather instantiated mono items grouped by def_id
+    let mut items_per_def_id: FxHashMap<_, Vec<_>> = Default::default();
+    for cgu in codegen_units {
+        for (&mono_item, _) in cgu.items() {
+            // Avoid variable-sized compiler-generated shims
+            if mono_item.is_user_defined() {
+                items_per_def_id.entry(mono_item.def_id()).or_default().push(mono_item);
+            }
+        }
+    }
+
+    // Output stats sorted by total instantiated size, from heaviest to lightest
+    let mut stats: Vec<_> = items_per_def_id
+        .into_iter()
+        .map(|(def_id, items)| {
+            let instantiation_count = items.len();
+            let size_estimate = items[0].size_estimate(tcx);
+            let total_estimate = instantiation_count * size_estimate;
+            (def_id, instantiation_count, size_estimate, total_estimate)
+        })
+        .collect();
+    stats.sort_unstable_by_key(|(_, _, _, total_estimate)| cmp::Reverse(*total_estimate));
+
+    if !stats.is_empty() {
+        writeln!(
+            file,
+            "| Item | Instantiation count | Estimated Cost Per Instantiation | Total Estimated Cost |"
+        )?;
+        writeln!(file, "| --- | ---: | ---: | ---: |")?;
+        for (def_id, instantiation_count, size_estimate, total_estimate) in stats {
+            let item = with_no_trimmed_paths!(tcx.def_path_str(def_id));
+            writeln!(
+                file,
+                "| {item} | {instantiation_count} | {size_estimate} | {total_estimate} |"
+            )?;
+        }
+    }
+
+    Ok(())
+}
+
 fn codegened_and_inlined_items<'tcx>(tcx: TyCtxt<'tcx>, (): ()) -> &'tcx DefIdSet {
     let (items, cgus) = tcx.collect_and_partition_mono_items(());
     let mut visited = DefIdSet::default();

compiler/rustc_session/src/options.rs

@@ -1294,6 +1294,9 @@ options! {
         computed `block` spans (one span encompassing a block's terminator and \
         all statements). If `-Z instrument-coverage` is also enabled, create \
         an additional `.html` file showing the computed coverage spans."),
+    dump_mono_stats: SwitchWithOptPath = (SwitchWithOptPath::Disabled,
+        parse_switch_with_opt_path, [UNTRACKED],
+        "output statistics about monomorphization collection (format: markdown)"),
     dwarf_version: Option<u32> = (None, parse_opt_number, [TRACKED],
         "version of DWARF debug information to emit (default: 2 or 4, depending on platform)"),
     dylib_lto: bool = (false, parse_bool, [UNTRACKED],

src/test/rustdoc-ui/z-help.stdout

@@ -35,6 +35,7 @@
     -Z            dump-mir-exclude-pass-number=val -- exclude the pass number when dumping MIR (used in tests) (default: no)
     -Z                       dump-mir-graphviz=val -- in addition to `.mir` files, create graphviz `.dot` files (and with `-Z instrument-coverage`, also create a `.dot` file for the MIR-derived coverage graph) (default: no)
     -Z                       dump-mir-spanview=val -- in addition to `.mir` files, create `.html` files to view spans for all `statement`s (including terminators), only `terminator` spans, or computed `block` spans (one span encompassing a block's terminator and all statements). If `-Z instrument-coverage` is also enabled, create an additional `.html` file showing the computed coverage spans.
+    -Z                         dump-mono-stats=val -- output statistics about monomorphization collection (format: markdown)
     -Z                           dwarf-version=val -- version of DWARF debug information to emit (default: 2 or 4, depending on platform)
     -Z                               dylib-lto=val -- enables LTO for dylib crate type
     -Z                        emit-stack-sizes=val -- emit a section containing stack size metadata (default: no)