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;
+ }
+ }
+ }
  _ => {}
  }
  }