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

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

r? ghost

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/mir/mod.rs

@@ -167,7 +167,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
     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

@@ -554,14 +554,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_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;
+                            }
+                        }
+                    }
                     _ => {}
                 }
             }

compiler/rustc_mir_transform/src/lib.rs

@@ -30,7 +30,7 @@ use rustc_middle::mir::{
     MirPhase, Operand, Place, ProjectionElem, Promoted, RuntimePhase, Rvalue, START_BLOCK,
     SourceInfo, Statement, StatementKind, TerminatorKind,
 };
-use rustc_middle::ty::{self, TyCtxt, TypeVisitableExt};
+use rustc_middle::ty::{self, Instance, TyCtxt, TypeVisitableExt};
 use rustc_middle::util::Providers;
 use rustc_middle::{bug, query, span_bug};
 use rustc_span::source_map::Spanned;
@@ -136,6 +136,7 @@ pub fn provide(providers: &mut Providers) {
         promoted_mir,
         deduced_param_attrs: deduce_param_attrs::deduced_param_attrs,
         coroutine_by_move_body_def_id: coroutine::coroutine_by_move_body_def_id,
+        build_codegen_mir,
         ..providers.queries
     };
 }
@@ -564,11 +565,11 @@ fn run_runtime_cleanup_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     }
 }
 
-fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-    fn o1<T>(x: T) -> WithMinOptLevel<T> {
-        WithMinOptLevel(1, x)
-    }
+fn o1<T>(x: T) -> WithMinOptLevel<T> {
+    WithMinOptLevel(1, x)
+}
 
+fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     // The main optimizations that we do on MIR.
     pm::run_passes(
         tcx,
@@ -609,7 +610,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &instsimplify::InstSimplify::AfterSimplifyCfg,
             &simplify::SimplifyLocals::BeforeConstProp,
             &dead_store_elimination::DeadStoreElimination::Initial,
-            &gvn::GVN,
+            &gvn::GVN::Polymorphic,
             &simplify::SimplifyLocals::AfterGVN,
             &dataflow_const_prop::DataflowConstProp,
             &single_use_consts::SingleUseConsts,
@@ -628,8 +629,6 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &multiple_return_terminators::MultipleReturnTerminators,
             &deduplicate_blocks::DeduplicateBlocks,
             &large_enums::EnumSizeOpt { discrepancy: 128 },
-            // Some cleanup necessary at least for LLVM and potentially other codegen backends.
-            &add_call_guards::CriticalCallEdges,
             // Cleanup for human readability, off by default.
             &prettify::ReorderBasicBlocks,
             &prettify::ReorderLocals,
@@ -689,6 +688,38 @@ fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
     body
 }
 
+pub fn build_codegen_mir<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Body<'tcx> {
+    let body = tcx.instance_mir(instance.def);
+    let mut body = instance.instantiate_mir_and_normalize_erasing_regions(
+        tcx,
+        ty::ParamEnv::reveal_all(),
+        ty::EarlyBinder::bind(body.clone()),
+    );
+    pm::run_passes(
+        tcx,
+        &mut body,
+        &[
+            // Validation calls layout::fn_can_unwind to figure out if a function can unwind, which
+            // always returns false if the current crate is compiled with -Cpanic=abort. So when
+            // a crate with panic=abort compiles MIR from a panic=unwind crate, we get validation
+            // failures. So we rely on the fact that validation only runs after passes? It's
+            // probably better to just delete that validation check.
+            &abort_unwinding_calls::AbortUnwindingCalls,
+            &gvn::GVN::PostMono,
+            // FIXME: Enabling this InstSimplify is required to fix the MIR from the
+            // unreachable_unchecked precondition check that UnreachablePropagation creates, but
+            // also enabling it breaks tests/codegen/issues/issue-122600-ptr-discriminant-update.rs
+            // LLVM appears to handle switches on i64 better than it handles icmp eq + br.
+            &instsimplify::InstSimplify::PostMono,
+            &o1(simplify_branches::SimplifyConstCondition::PostMono),
+            &o1(simplify::SimplifyCfg::PostMono),
+            &add_call_guards::CriticalCallEdges,
+        ],
+        Some(MirPhase::Runtime(RuntimePhase::Codegen)),
+    );
+    body
+}
+
 /// Fetch all the promoteds of an item and prepare their MIR bodies to be ready for
 /// constant evaluation once all generic parameters become known.
 fn promoted_mir(tcx: TyCtxt<'_>, def: LocalDefId) -> &IndexVec<Promoted, Body<'_>> {

compiler/rustc_mir_transform/src/pass_manager.rs

@@ -252,7 +252,10 @@ fn run_passes_inner<'tcx>(
 
         let validate =
             (validate_each & tcx.sess.opts.unstable_opts.validate_mir & !body.should_skip())
-                || new_phase == MirPhase::Runtime(RuntimePhase::Optimized);
+                || matches!(
+                    new_phase,
+                    MirPhase::Runtime(RuntimePhase::Optimized | RuntimePhase::Codegen)
+                );
         let lint = tcx.sess.opts.unstable_opts.lint_mir & !body.should_skip();
         if validate {
             validate_body(tcx, body, format!("after phase change to {}", new_phase.name()));
@@ -265,7 +268,7 @@ fn run_passes_inner<'tcx>(
     }
 }
 
-pub(super) fn validate_body<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, when: String) {
+fn validate_body<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, when: String) {
     validate::Validator { when, mir_phase: body.phase }.run_pass(tcx, body);
 }
 

compiler/rustc_mir_transform/src/simplify.rs

@@ -47,6 +47,7 @@ pub(super) enum SimplifyCfg {
     Final,
     MakeShim,
     AfterUnreachableEnumBranching,
+    PostMono,
 }
 
 impl SimplifyCfg {
@@ -62,6 +63,7 @@ impl SimplifyCfg {
             SimplifyCfg::AfterUnreachableEnumBranching => {
                 "SimplifyCfg-after-unreachable-enum-branching"
             }
+            SimplifyCfg::PostMono => "SimplifyCfg-post-mono",
         }
     }
 }

compiler/rustc_mir_transform/src/simplify_branches.rs

@@ -5,6 +5,7 @@ use tracing::trace;
 pub(super) enum SimplifyConstCondition {
     AfterConstProp,
     Final,
+    PostMono,
 }
 
 /// A pass that replaces a branch with a goto when its condition is known.
@@ -13,6 +14,7 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyConstCondition {
         match self {
             SimplifyConstCondition::AfterConstProp => "SimplifyConstCondition-after-const-prop",
             SimplifyConstCondition::Final => "SimplifyConstCondition-final",
+            SimplifyConstCondition::PostMono => "SimplifyConstCondition-post-mono",
         }
     }
 

compiler/rustc_mir_transform/src/validate.rs

@@ -395,7 +395,7 @@ impl<'a, 'tcx> Visitor<'tcx> for CfgChecker<'a, 'tcx> {
                     // the return edge from the call. FIXME(tmiasko): Since this is a strictly code
                     // generation concern, the code generation should be responsible for handling
                     // it.
-                    if self.mir_phase >= MirPhase::Runtime(RuntimePhase::Optimized)
+                    if self.mir_phase >= MirPhase::Runtime(RuntimePhase::Codegen)
                         && self.is_critical_call_edge(target, unwind)
                     {
                         self.fail(

compiler/rustc_monomorphize/src/collector.rs

@@ -1207,7 +1207,7 @@ fn collect_items_of_instance<'tcx>(
     mentioned_items: &mut MonoItems<'tcx>,
     mode: CollectionMode,
 ) {
-    let body = tcx.instance_mir(instance.def);
+    let body = tcx.build_codegen_mir(instance);
     // Naively, in "used" collection mode, all functions get added to *both* `used_items` and
     // `mentioned_items`. Mentioned items processing will then notice that they have already been
     // visited, but at that point each mentioned item has been monomorphized, added to the

tests/codegen/issues/issue-122600-ptr-discriminant-update.rs

@@ -1,4 +1,4 @@
-//@ compile-flags: -O
+//@ compile-flags: -O -Zmir-enable-passes=-InstSimplify-post-mono
 //@ min-llvm-version: 19
 
 #![crate_type = "lib"]