Auto merge of #131650 - saethlin:post-mono-mir-opts, r=<try>

Add post-mono MIR passes to make mono-reachable analysis more accurate

As of #131650 (comment) I believe most of the incr overhead comes from re-computing, re-encoding, and loading a lot more MIR when all we're actually doing is traversing through it. I think that can be addressed by caching a query that looks up the mentioned/used items for an Instance.

I think the full-build regressions are pretty much just the expense of cloning, then monomorphizing, then caching the MIR.

Author: bors

compiler/rustc_codegen_cranelift/src/base.rs

@@ -41,7 +41,7 @@ pub(crate) fn codegen_fn<'tcx>(
     let symbol_name = tcx.symbol_name(instance).name.to_string();
     let _timer = tcx.prof.generic_activity_with_arg("codegen fn", &*symbol_name);
 
-    let mir = tcx.instance_mir(instance.def);
+    let mir = tcx.build_codegen_mir(instance);
     let _mir_guard = crate::PrintOnPanic(|| {
         let mut buf = Vec::new();
         with_no_trimmed_paths!({

compiler/rustc_codegen_ssa/src/base.rs

@@ -416,7 +416,7 @@ pub(crate) fn codegen_instance<'a, 'tcx: 'a, Bx: BuilderMethods<'a, 'tcx>>(
     // release builds.
     info!("codegen_instance({})", instance);
 
-    mir::codegen_mir::<Bx>(cx, instance);
+    mir::lower_mir::<Bx>(cx, instance);
 }
 
 /// Creates the `main` function which will initialize the rust runtime and call

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -159,15 +159,15 @@ impl<'tcx, V: CodegenObject> LocalRef<'tcx, V> {
 ///////////////////////////////////////////////////////////////////////////
 
 #[instrument(level = "debug", skip(cx))]
-pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+pub fn lower_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     cx: &'a Bx::CodegenCx,
     instance: Instance<'tcx>,
 ) {
     assert!(!instance.args.has_infer());
 
     let llfn = cx.get_fn(instance);
 
-    let mir = cx.tcx().instance_mir(instance.def);
+    let mir = cx.tcx().build_codegen_mir(instance);
 
     let fn_abi = cx.fn_abi_of_instance(instance, ty::List::empty());
     debug!("fn_abi: {:?}", fn_abi);

compiler/rustc_middle/src/mir/mod.rs

@@ -159,6 +159,7 @@ impl RuntimePhase {
             "initial" => Self::Initial,
             "post_cleanup" | "post-cleanup" | "postcleanup" => Self::PostCleanup,
             "optimized" => Self::Optimized,
+            "codegen" => Self::Codegen,
             _ => bug!("Unknown runtime phase: '{}'", phase),
         }
     }

compiler/rustc_middle/src/mir/mono.rs

@@ -73,7 +73,7 @@ impl<'tcx> MonoItem<'tcx> {
                     InstanceKind::Item(..)
                     | InstanceKind::DropGlue(..)
                     | InstanceKind::AsyncDropGlueCtorShim(..) => {
-                        let mir = tcx.instance_mir(instance.def);
+                        let mir = tcx.build_codegen_mir(instance);
                         mir.basic_blocks.iter().map(|bb| bb.statements.len() + 1).sum()
                     }
                     // Other compiler-generated shims size estimate: 1

compiler/rustc_middle/src/mir/syntax.rs

@@ -99,6 +99,7 @@ impl MirPhase {
             MirPhase::Runtime(RuntimePhase::Initial) => "runtime",
             MirPhase::Runtime(RuntimePhase::PostCleanup) => "runtime-post-cleanup",
             MirPhase::Runtime(RuntimePhase::Optimized) => "runtime-optimized",
+            MirPhase::Runtime(RuntimePhase::Codegen) => "codegen",
         }
     }
 
@@ -153,6 +154,7 @@ pub enum RuntimePhase {
     /// * [`ProjectionElem::Deref`] of `Box`
     PostCleanup = 1,
     Optimized = 2,
+    Codegen = 3,
 }
 
 ///////////////////////////////////////////////////////////////////////////

compiler/rustc_middle/src/query/mod.rs

@@ -555,14 +555,22 @@ rustc_queries! {
         desc { |tcx| "verify auto trait bounds for coroutine interior type `{}`", tcx.def_path_str(key) }
     }
 
-    /// MIR after our optimization passes have run. This is MIR that is ready
-    /// for codegen. This is also the only query that can fetch non-local MIR, at present.
+    /// Polymorphic MIR after our pre-mono optimization passes have run. This is the MIR that
+    /// crates export.
     query optimized_mir(key: DefId) -> &'tcx mir::Body<'tcx> {
         desc { |tcx| "optimizing MIR for `{}`", tcx.def_path_str(key) }
         cache_on_disk_if { key.is_local() }
         separate_provide_extern
     }
 
+    /// MIR for a specific Instance ready for codegen. This is `optimized_mir` but monomorphized
+    /// and with extra transforms applied.
+    query build_codegen_mir(key: ty::Instance<'tcx>) -> &'tcx mir::Body<'tcx> {
+        desc { |tcx| "finalizing codegen MIR for `{}`", tcx.def_path_str_with_args(key.def_id(), key.args) }
+        cache_on_disk_if { true }
+        arena_cache
+    }
+
     /// Checks for the nearest `#[coverage(off)]` or `#[coverage(on)]` on
     /// this def and any enclosing defs, up to the crate root.
     ///

compiler/rustc_middle/src/ty/mod.rs

@@ -1718,6 +1718,14 @@ impl<'tcx> TyCtxt<'tcx> {
         }
     }
 
+    pub fn codegen_mir(self, instance: ty::Instance<'tcx>) -> &'tcx Body<'tcx> {
+        if self.sess.opts.incremental == None {
+            self.build_codegen_mir(instance)
+        } else {
+            self.instance_mir(instance.def)
+        }
+    }
+
     /// Returns the possibly-auto-generated MIR of a [`ty::InstanceKind`].
     #[instrument(skip(self), level = "debug")]
     pub fn instance_mir(self, instance: ty::InstanceKind<'tcx>) -> &'tcx Body<'tcx> {

compiler/rustc_mir_transform/src/gvn.rs

@@ -85,6 +85,7 @@
 use std::borrow::Cow;
 
 use either::Either;
+use rustc_ast::attr;
 use rustc_const_eval::const_eval::DummyMachine;
 use rustc_const_eval::interpret::{
     ImmTy, Immediate, InterpCx, MemPlaceMeta, MemoryKind, OpTy, Projectable, Scalar,
@@ -101,17 +102,27 @@ use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::layout::{HasParamEnv, LayoutOf};
 use rustc_middle::ty::{self, Ty, TyCtxt};
-use rustc_span::DUMMY_SP;
 use rustc_span::def_id::DefId;
+use rustc_span::{DUMMY_SP, sym};
 use rustc_target::abi::{self, Abi, FIRST_VARIANT, FieldIdx, Primitive, Size, VariantIdx};
 use smallvec::SmallVec;
 use tracing::{debug, instrument, trace};
 
 use crate::ssa::{AssignedValue, SsaLocals};
 
-pub(super) struct GVN;
+pub(super) enum GVN {
+    Polymorphic,
+    PostMono,
+}
 
 impl<'tcx> crate::MirPass<'tcx> for GVN {
+    fn name(&self) -> &'static str {
+        match self {
+            GVN::Polymorphic => "GVN",
+            GVN::PostMono => "GVN-post-mono",
+        }
+    }
+
     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
         sess.mir_opt_level() >= 2
     }
@@ -125,7 +136,22 @@ impl<'tcx> crate::MirPass<'tcx> for GVN {
         // Clone dominators because we need them while mutating the body.
         let dominators = body.basic_blocks.dominators().clone();
 
-        let mut state = VnState::new(tcx, body, param_env, &ssa, dominators, &body.local_decls);
+        let preserve_ub_checks = match self {
+            GVN::Polymorphic => {
+                attr::contains_name(tcx.hir().krate_attrs(), sym::rustc_preserve_ub_checks)
+            }
+            GVN::PostMono => false,
+        };
+
+        let mut state = VnState::new(
+            tcx,
+            body,
+            param_env,
+            &ssa,
+            dominators,
+            &body.local_decls,
+            preserve_ub_checks,
+        );
         ssa.for_each_assignment_mut(
             body.basic_blocks.as_mut_preserves_cfg(),
             |local, value, location| {
@@ -260,6 +286,7 @@ struct VnState<'body, 'tcx> {
     ssa: &'body SsaLocals,
     dominators: Dominators<BasicBlock>,
     reused_locals: BitSet<Local>,
+    preserve_ub_checks: bool,
 }
 
 impl<'body, 'tcx> VnState<'body, 'tcx> {
@@ -270,6 +297,7 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
         ssa: &'body SsaLocals,
         dominators: Dominators<BasicBlock>,
         local_decls: &'body LocalDecls<'tcx>,
+        preserve_ub_checks: bool,
     ) -> Self {
         // Compute a rough estimate of the number of values in the body from the number of
         // statements. This is meant to reduce the number of allocations, but it's all right if
@@ -292,6 +320,7 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
             ssa,
             dominators,
             reused_locals: BitSet::new_empty(local_decls.len()),
+            preserve_ub_checks,
         }
     }
 
@@ -530,7 +559,14 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
                         .tcx
                         .offset_of_subfield(self.ecx.param_env(), layout, fields.iter())
                         .bytes(),
-                    NullOp::UbChecks => return None,
+                    NullOp::UbChecks => {
+                        if self.preserve_ub_checks {
+                            return None;
+                        } else {
+                            let val = ImmTy::from_bool(self.tcx.sess.ub_checks(), self.tcx);
+                            return Some(val.into());
+                        }
+                    }
                 };
                 let usize_layout = self.ecx.layout_of(self.tcx.types.usize).unwrap();
                 let imm = ImmTy::from_uint(val, usize_layout);

compiler/rustc_mir_transform/src/instsimplify.rs

@@ -16,13 +16,15 @@ use crate::take_array;
 pub(super) enum InstSimplify {
     BeforeInline,
     AfterSimplifyCfg,
+    PostMono,
 }
 
 impl<'tcx> crate::MirPass<'tcx> for InstSimplify {
     fn name(&self) -> &'static str {
         match self {
             InstSimplify::BeforeInline => "InstSimplify-before-inline",
             InstSimplify::AfterSimplifyCfg => "InstSimplify-after-simplifycfg",
+            InstSimplify::PostMono => "InstSimplify-post-mono",
         }
     }
 
@@ -51,6 +53,29 @@ impl<'tcx> crate::MirPass<'tcx> for InstSimplify {
                         ctx.simplify_ptr_aggregate(&statement.source_info, rvalue);
                         ctx.simplify_cast(rvalue);
                     }
+                    StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(ref op)) => {
+                        // UnreachablePropagation likes to generate this MIR:
+                        //
+                        // _1 = UbChecks();
+                        // assume(copy _1);
+                        // _2 = unreachable_unchecked::precondition_check() -> [return: bb2, unwind unreachable];
+                        //
+                        // Which is mind-bending but correct. When UbChecks is false, we
+                        // assume(false) which is unreachable, and we never hit the precondition
+                        // check. When UbChecks is true, we assume(true) and fall through to the
+                        // precondition check.
+                        //
+                        // So the branch on UbChecks is implicit, which is both clever and makes
+                        // the rest of MIR optimizations unable to delete this precondition check
+                        // call when UB checks are off.
+                        if let Some(ConstOperand { const_, .. }) = op.constant() {
+                            if let Some(false) = const_.try_to_bool() {
+                                block.statements.clear();
+                                block.terminator_mut().kind = TerminatorKind::Unreachable;
+                                break;
+                            }
+                        }
+                    }
                     _ => {}
                 }
             }