Outline and deduplicate cold panic and drop paths in MIR and LLVM

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1563,6 +1563,36 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         let cold_inline = llvm::AttributeKind::Cold.create_attr(self.llcx);
         attributes::apply_to_callsite(llret, llvm::AttributePlace::Function, &[cold_inline]);
     }
+
+    fn apply_attrs_to_cold_callsite(&mut self, llret: &'ll Value, noinline: bool) {
+        unsafe extern "C" {
+            fn LLVMRustEnableColdOutline();
+        }
+        if self
+            .tcx
+            .sess
+            .opts
+            .unstable_opts
+            .mir_enable_passes
+            .iter()
+            .rev()
+            .find(|(name, _)| name == "ColdBlockOutliner")
+            .is_some_and(|(_, enabled)| *enabled)
+        {
+            unsafe { LLVMRustEnableColdOutline() };
+        }
+        let cold_attr = llvm::AttributeKind::Cold.create_attr(self.llcx);
+        if noinline {
+            let noinline_attr = llvm::AttributeKind::NoInline.create_attr(self.llcx);
+            attributes::apply_to_callsite(
+                llret,
+                llvm::AttributePlace::Function,
+                &[cold_attr, noinline_attr],
+            );
+        } else {
+            attributes::apply_to_callsite(llret, llvm::AttributePlace::Function, &[cold_attr]);
+        }
+    }
 }
 
 impl<'ll> StaticBuilderMethods for Builder<'_, 'll, '_> {

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -267,6 +267,26 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
             return MergingSucc::False;
         }
 
+        let is_cleanup = fx.mir[self.bb].is_cleanup;
+        let is_cold = fx.cold_blocks[self.bb];
+        let apply_callsite_attrs = |bx: &mut Bx, call_llval| {
+            if is_cleanup {
+                bx.apply_attrs_to_cleanup_callsite(call_llval);
+            } else if is_cold && super::should_mark_cold_callsites(bx.tcx()) {
+                let tcx = bx.tcx();
+                let noinline = instance.is_some_and(|inst| {
+                    matches!(inst.def, ty::InstanceKind::Shim(ty::ShimKind::DropGlue(..)))
+                        || (super::should_outline_cold_calls(tcx)
+                            && (tcx
+                                .codegen_instance_attrs(inst.def)
+                                .flags
+                                .contains(super::CodegenFnAttrFlags::COLD)
+                                || super::is_try_from_residual(tcx, inst.def_id())))
+                });
+                bx.apply_attrs_to_cold_callsite(call_llval, noinline);
+            }
+        };
+
         if let Some(unwind_block) = unwind_block {
             let ret_llbb = if let Some((_, target)) = destination {
                 self.llbb_with_cleanup(fx, target)
@@ -285,9 +305,7 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
                 self.funclet(fx),
                 instance,
             );
-            if fx.mir[self.bb].is_cleanup {
-                bx.apply_attrs_to_cleanup_callsite(invokeret);
-            }
+            apply_callsite_attrs(bx, invokeret);
 
             if let Some((ret_dest, target)) = destination {
                 bx.switch_to_block(fx.llbb(target));
@@ -315,9 +333,7 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
                 self.funclet(fx),
                 instance,
             );
-            if fx.mir[self.bb].is_cleanup {
-                bx.apply_attrs_to_cleanup_callsite(llret);
-            }
+            apply_callsite_attrs(bx, llret);
 
             if let Some((ret_dest, target)) = destination {
                 for &(tmp, size) in lifetime_ends_after_call {

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -1,13 +1,15 @@
 use std::iter;
+use std::sync::LazyLock;
 
+use rustc_attr_ir::lang_items::LangItem;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir;
 use rustc_middle::mir::{Body, Local, UnwindTerminateReason, traversal};
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, TyAndLayout};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt, TypeFoldable, TypeVisitableExt};
-use rustc_span::{ErrorGuaranteed, bug, span_bug};
+use rustc_span::{ErrorGuaranteed, bug, span_bug, sym};
 use rustc_target::callconv::{FnAbi, IndirectMode, PassMode};
 use tracing::{debug, instrument};
 
@@ -665,6 +667,98 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     args
 }
 
+static COLD_CATCH_UNWIND: LazyLock<bool> =
+    LazyLock::new(|| std::env::var("RUSTC_EXP_COLD_CATCH_UNWIND").is_ok_and(|v| v == "1"));
+
+static COLD_OUTLINE: LazyLock<bool> =
+    LazyLock::new(|| std::env::var("RUSTC_EXP_COLD_OUTLINE").is_ok_and(|v| v == "1"));
+
+pub(super) fn should_outline_cold_calls(tcx: TyCtxt<'_>) -> bool {
+    *COLD_OUTLINE
+        || tcx
+            .sess
+            .opts
+            .unstable_opts
+            .mir_enable_passes
+            .iter()
+            .rev()
+            .find(|(name, _)| name == "ColdBlockOutliner")
+            .is_some_and(|(_, enabled)| *enabled)
+}
+
+pub(super) fn should_mark_cold_callsites(tcx: TyCtxt<'_>) -> bool {
+    *COLD_CATCH_UNWIND
+        || should_outline_cold_calls(tcx)
+        || tcx
+            .sess
+            .opts
+            .unstable_opts
+            .mir_enable_passes
+            .iter()
+            .rev()
+            .find(|(name, _)| name == "KnownVariantDrops")
+            .is_some_and(|(_, enabled)| *enabled)
+}
+
+pub(super) fn is_try_from_residual(tcx: TyCtxt<'_>, def_id: rustc_hir::def_id::DefId) -> bool {
+    tcx.is_lang_item(def_id, LangItem::TryTraitFromResidual)
+        || tcx
+            .trait_item_of(def_id)
+            .is_some_and(|ti| tcx.is_lang_item(ti, LangItem::TryTraitFromResidual))
+}
+
+/// Marks the successors of `switchInt(catch_unwind(..))` that are taken for a nonzero result
+/// (a caught panic) as cold.
+fn mark_caught_panic_blocks_cold<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    mir: &mir::Body<'tcx>,
+    cold_blocks: &mut IndexVec<mir::BasicBlock, bool>,
+) {
+    for bb_data in mir.basic_blocks.iter() {
+        let mir::TerminatorKind::Call { ref func, destination, target: Some(target), .. } =
+            bb_data.terminator().kind
+        else {
+            continue;
+        };
+        let ty::FnDef(def_id, _) = *func.ty(&mir.local_decls, tcx).kind() else { continue };
+        if !tcx.is_intrinsic(def_id, sym::catch_unwind) {
+            continue;
+        }
+        let Some(result) = destination.as_local() else { continue };
+        let target_data = &mir.basic_blocks[target];
+        let mir::TerminatorKind::SwitchInt { ref discr, ref targets } =
+            target_data.terminator().kind
+        else {
+            continue;
+        };
+        if discr.place().and_then(|place| place.as_local()) != Some(result) {
+            continue;
+        }
+        let result_rewritten =
+            target_data.statements.iter().any(|statement| match &statement.kind {
+                mir::StatementKind::Assign(assign) => assign.0.local == result,
+                mir::StatementKind::SetDiscriminant { place, .. } => place.local == result,
+                mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) => {
+                    *local == result
+                }
+                _ => false,
+            });
+        if result_rewritten {
+            continue;
+        }
+        let no_panic = targets.target_for_value(0);
+        for (value, bb) in targets.iter() {
+            if value != 0 && bb != no_panic {
+                cold_blocks[bb] = true;
+            }
+        }
+        let otherwise = targets.otherwise();
+        if targets.iter().any(|(value, _)| value == 0) && otherwise != no_panic {
+            cold_blocks[otherwise] = true;
+        }
+    }
+}
+
 fn find_cold_blocks<'tcx>(
     tcx: TyCtxt<'tcx>,
     mir: &mir::Body<'tcx>,
@@ -674,6 +768,12 @@ fn find_cold_blocks<'tcx>(
     let mut cold_blocks: IndexVec<mir::BasicBlock, bool> =
         IndexVec::from_elem(false, &mir.basic_blocks);
 
+    if should_mark_cold_callsites(tcx) {
+        mark_caught_panic_blocks_cold(tcx, mir, &mut cold_blocks);
+    }
+
+    let mark_residual_cold = should_outline_cold_calls(tcx);
+
     // Traverse all basic blocks from end of the function to the start.
     for (bb, bb_data) in traversal::postorder(mir) {
         let terminator = bb_data.terminator();
@@ -684,7 +784,8 @@ fn find_cold_blocks<'tcx>(
             | mir::TerminatorKind::TailCall { ref func, .. }
                 if let ty::FnDef(def_id, ..) = *func.ty(local_decls, tcx).kind()
                     && let attrs = tcx.codegen_fn_attrs(def_id)
-                    && attrs.flags.contains(CodegenFnAttrFlags::COLD) =>
+                    && (attrs.flags.contains(CodegenFnAttrFlags::COLD)
+                        || (mark_residual_cold && is_try_from_residual(tcx, def_id))) =>
             {
                 cold_blocks[bb] = true;
                 continue;
@@ -699,8 +800,10 @@ fn find_cold_blocks<'tcx>(
             _ => {}
         }
 
-        // If all successors of a BB are cold and there's at least one of them, mark this BB as cold
-        let mut succ = terminator.successors();
+        // If all non-cleanup successors of a BB are cold and there's at least one of them, mark this BB as cold
+        let mut succ = terminator
+            .successors()
+            .filter(|&s| !mark_residual_cold || !mir.basic_blocks[s].is_cleanup);
         if let Some(first) = succ.next()
             && cold_blocks[first]
             && succ.all(|s| cold_blocks[s])

compiler/rustc_codegen_ssa/src/traits/builder.rs

@@ -690,4 +690,8 @@ pub trait BuilderMethods<'a, 'tcx>:
     fn zext(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value;
 
     fn apply_attrs_to_cleanup_callsite(&mut self, llret: Self::Value);
+
+    fn apply_attrs_to_cold_callsite(&mut self, llret: Self::Value, _noinline: bool) {
+        self.apply_attrs_to_cleanup_callsite(llret);
+    }
 }