diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs index 279c11257a6f9..085c96d98a787 100644 --- a/compiler/rustc_codegen_llvm/src/builder.rs +++ b/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, '_> { diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs index c1c690791f195..2ae6947d95d02 100644 --- a/compiler/rustc_codegen_ssa/src/mir/block.rs +++ b/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 { diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs index a6ad7282ca08d..c80ce1536381b 100644 --- a/compiler/rustc_codegen_ssa/src/mir/mod.rs +++ b/compiler/rustc_codegen_ssa/src/mir/mod.rs @@ -1,5 +1,7 @@ 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; @@ -7,7 +9,7 @@ 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 = + LazyLock::new(|| std::env::var("RUSTC_EXP_COLD_CATCH_UNWIND").is_ok_and(|v| v == "1")); + +static COLD_OUTLINE: LazyLock = + 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, +) { + 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 = 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]) diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs index 203274ccbb9fb..27002452e62c2 100644 --- a/compiler/rustc_codegen_ssa/src/traits/builder.rs +++ b/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); + } } diff --git a/compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp index fbb810d5423b2..3fd2ce5b8e1da 100644 --- a/compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp +++ b/compiler/rustc_llvm/llvm-wrapper/PassWrapper.cpp @@ -12,6 +12,7 @@ #include "llvm/CodeGen/CommandFlags.h" #include "llvm/IR/AssemblyAnnotationWriter.h" #include "llvm/IR/AutoUpgrade.h" +#include "llvm/IR/EHPersonalities.h" #include "llvm/IR/LegacyPassManager.h" #include "llvm/IR/PassManager.h" #include "llvm/IR/Verifier.h" @@ -34,6 +35,7 @@ #include "llvm/Transforms/IPO/FunctionImport.h" #include "llvm/Transforms/IPO/Internalize.h" #include "llvm/Transforms/IPO/LowerTypeTests.h" +#include "llvm/Transforms/IPO/MergeFunctions.h" #include "llvm/Transforms/IPO/ThinLTOBitcodeWriter.h" #include "llvm/Transforms/Instrumentation/AddressSanitizer.h" #include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h" @@ -43,12 +45,15 @@ #include "llvm/Transforms/Instrumentation/RealtimeSanitizer.h" #include "llvm/Transforms/Instrumentation/ThreadSanitizer.h" #include "llvm/Transforms/Scalar/AnnotationRemarks.h" +#include "llvm/Transforms/Utils/CodeExtractor.h" #if LLVM_VERSION_GE(23, 0) #include "llvm/Transforms/Utils/AssignGUID.h" #endif #include "llvm/Transforms/Utils/CanonicalizeAliases.h" #include "llvm/Transforms/Utils/FunctionImportUtils.h" #include "llvm/Transforms/Utils/NameAnonGlobals.h" +#include +#include #include #include #include @@ -58,6 +63,385 @@ using namespace llvm; +static std::atomic GEnableColdOutline{false}; + +extern "C" void LLVMRustEnableColdOutline() { + GEnableColdOutline.store(true, std::memory_order_relaxed); +} + +namespace { +struct ColdInvokeOutlinerPass + : public RequiredPassInfoMixin { + PreservedAnalyses run(Module &M, ModuleAnalysisManager &) { + bool Changed = false; + SmallVector WorkFuncs; + for (Function &F : M) { + if (!F.isDeclaration()) + WorkFuncs.push_back(&F); + } + for (Function *FPtr : WorkFuncs) { + Function &F = *FPtr; + if (F.hasOptNone() || F.hasFnAttribute(Attribute::Cold) || + F.hasFnAttribute(Attribute::NoReturn) || + F.hasFnAttribute(Attribute::Naked)) + continue; + if (F.hasPersonalityFn() && + isScopedEHPersonality(classifyEHPersonality(F.getPersonalityFn()))) + continue; + + auto collectColdOnlyAllocas = + [&](BasicBlock *TargetBB, SmallVectorImpl &ColdAllocas, + SmallVectorImpl &DeadLifetimes) { + for (Instruction &EntryI : F.getEntryBlock()) { + auto *AI = dyn_cast(&EntryI); + if (!AI || !AI->isStaticAlloca()) + continue; + bool OnlyInTarget = true; + bool UsedInTarget = false; + SmallVector Worklist(AI->users()); + SmallPtrSet Visited; + SmallVector LocalLifetimes; + while (!Worklist.empty()) { + User *U = Worklist.pop_back_val(); + if (!Visited.insert(U).second) + continue; + auto *UI = dyn_cast(U); + if (!UI) { + OnlyInTarget = false; + break; + } + if (auto *Intr = dyn_cast(UI)) { + if (Intr->getIntrinsicID() == Intrinsic::lifetime_start || + Intr->getIntrinsicID() == Intrinsic::lifetime_end) { + if (Intr->getParent() != TargetBB) + LocalLifetimes.push_back(Intr); + continue; + } + } + if (UI->getParent() == TargetBB) { + UsedInTarget = true; + if (isa(UI) || isa(UI)) { + for (User *SubU : UI->users()) + Worklist.push_back(SubU); + } + } else if (isa(UI) || isa(UI)) { + for (User *SubU : UI->users()) + Worklist.push_back(SubU); + } else { + OnlyInTarget = false; + break; + } + } + if (OnlyInTarget && UsedInTarget) { + ColdAllocas.push_back(AI); + DeadLifetimes.append(LocalLifetimes.begin(), + LocalLifetimes.end()); + } + } + }; + + auto hoistAllocasToEntry = [](Function *OutF) { + BasicBlock &Entry = OutF->getEntryBlock(); + SmallVector AllocasToMove; + for (BasicBlock &OutBB : *OutF) { + if (&OutBB == &Entry) + continue; + for (Instruction &OutI : OutBB) { + if (auto *AI = dyn_cast(&OutI)) + AllocasToMove.push_back(AI); + } + } + for (AllocaInst *AI : AllocasToMove) { + AI->moveBefore(Entry.getFirstInsertionPt()); + } + }; + + SmallVector Candidates; + for (BasicBlock &BB : F) { + if (&BB == &F.getEntryBlock() || BB.isEHPad() || BB.hasAddressTaken() || + pred_empty(&BB)) + continue; + + if (auto *II = dyn_cast(BB.getTerminator())) { + if (BB.size() < 2) + continue; + if (!II->hasFnAttr(Attribute::Cold) && !II->doesNotReturn()) + continue; + if (Function *Callee = II->getCalledFunction()) { + if (Callee->getName().contains(".cold")) + continue; + } + BasicBlock *NormalDest = II->getNormalDest(); + if (!NormalDest->phis().empty() || + !isa(NormalDest->getFirstNonPHIIt())) + continue; + BasicBlock *UnwindDest = II->getUnwindDest(); + bool Valid = true; + for (PHINode &PN : UnwindDest->phis()) { + int Idx = PN.getBasicBlockIndex(&BB); + if (Idx >= 0) { + if (auto *DefI = + dyn_cast(PN.getIncomingValue(Idx))) { + if (DefI->getParent() == &BB) { + Valid = false; + break; + } + } + } + } + for (Instruction &I : BB) { + if (!Valid) + break; + if (isa(&I) || I.getType()->isTokenTy()) { + Valid = false; + break; + } + if (&I != II) { + for (User *U : I.users()) { + auto *UI = dyn_cast(U); + if (!UI || UI->getParent() != &BB) { + Valid = false; + break; + } + } + } + } + if (Valid) + Candidates.push_back(&BB); + } else if (isa(BB.getTerminator())) { + CallInst *ColdCI = nullptr; + bool Valid = true; + for (Instruction &I : BB) { + if (isa(&I) || I.getType()->isTokenTy()) { + Valid = false; + break; + } + if (auto *CI = dyn_cast(&I)) { + if (CI->hasFnAttr(Attribute::Cold)) { + ColdCI = CI; + if (Function *Callee = CI->getCalledFunction()) { + if (Callee->getName().contains(".cold")) { + Valid = false; + break; + } + } + } + } + for (User *U : I.users()) { + auto *UI = dyn_cast(U); + if (!UI || UI->getParent() != &BB) { + Valid = false; + break; + } + } + if (!Valid) + break; + } + if (Valid && ColdCI) + Candidates.push_back(&BB); + } + } + + unsigned ColdIdx = 0; + for (BasicBlock *BB : Candidates) { + auto *OrigII = dyn_cast(BB->getTerminator()); + if (!OrigII && !isa(BB->getTerminator())) + continue; + + BasicBlock *NormalDest = nullptr; + BasicBlock *UnwindDest = nullptr; + SmallVector, 4> SavedPhis; + SmallVector InnerArgs; + SmallVector Bundles; + CallInst *InnerCI = nullptr; + UnreachableInst *TempUnreach = nullptr; + DebugLoc OrigDL; + + CallBase *ColdCall = nullptr; + unsigned NonTrivialInsts = 0; + for (Instruction &I : *BB) { + if (isa(&I)) + continue; + if (auto *Intr = dyn_cast(&I)) { + if (Intr->getIntrinsicID() == Intrinsic::lifetime_start || + Intr->getIntrinsicID() == Intrinsic::lifetime_end) + continue; + } + ++NonTrivialInsts; + if (auto *CB = dyn_cast(&I)) { + if (CB->hasFnAttr(Attribute::Cold) || CB == OrigII) + ColdCall = CB; + } + } + if (!ColdCall) + continue; + unsigned OrigCallArgSize = ColdCall->arg_size(); + if (NonTrivialInsts < 2) { + bool HasUniqueGlobalArg = false; + for (Value *Arg : ColdCall->args()) { + if (auto *GV = dyn_cast(Arg->stripPointerCasts())) { + if (GV->hasOneUse()) { + HasUniqueGlobalArg = true; + break; + } + } + } + if (HasUniqueGlobalArg) + continue; + } + + if (OrigII) { + NormalDest = OrigII->getNormalDest(); + UnwindDest = OrigII->getUnwindDest(); + for (PHINode &PN : UnwindDest->phis()) { + int Idx = PN.getBasicBlockIndex(BB); + if (Idx >= 0) { + SavedPhis.emplace_back(&PN, PN.getIncomingValue(Idx)); + PN.removeIncomingValue(Idx, /*DeletePHIIfEmpty=*/false); + } + } + InnerArgs.assign(OrigII->arg_begin(), OrigII->arg_end()); + OrigII->getOperandBundlesAsDefs(Bundles); + InnerCI = CallInst::Create( + OrigII->getFunctionType(), OrigII->getCalledOperand(), InnerArgs, + Bundles, OrigII->getName(), OrigII->getIterator()); + InnerCI->setCallingConv(OrigII->getCallingConv()); + InnerCI->setAttributes(OrigII->getAttributes()); + InnerCI->setDebugLoc(OrigII->getDebugLoc()); + TempUnreach = + new UnreachableInst(F.getContext(), OrigII->getIterator()); + OrigDL = OrigII->getDebugLoc(); + OrigII->eraseFromParent(); + } + + SmallVector ColdAllocas; + SmallVector DeadLifetimes; + collectColdOnlyAllocas(BB, ColdAllocas, DeadLifetimes); + for (Instruction *DeadI : DeadLifetimes) + DeadI->eraseFromParent(); + for (AllocaInst *AI : ColdAllocas) + AI->moveBefore(BB->getFirstNonPHIIt()); + + SmallPtrSet ColdAllocaSet(ColdAllocas.begin(), + ColdAllocas.end()); + auto getBaseColdAlloca = [&](Value *V) -> AllocaInst * { + while (true) { + if (auto *AI = dyn_cast(V)) + return ColdAllocaSet.contains(AI) ? AI : nullptr; + if (auto *GEP = dyn_cast(V)) { + V = GEP->getPointerOperand(); + continue; + } + if (auto *BC = dyn_cast(V)) { + V = BC->getOperand(0); + continue; + } + return nullptr; + } + }; + SmallPtrSet ReferencedColdAllocas; + for (Instruction &I : *BB) { + if (auto *SI = dyn_cast(&I)) { + AllocaInst *ValAI = getBaseColdAlloca(SI->getValueOperand()); + AllocaInst *PtrAI = getBaseColdAlloca(SI->getPointerOperand()); + if (ValAI && PtrAI && ValAI != PtrAI) + ReferencedColdAllocas.insert(ValAI); + } + } + SmallVector ProxiedConsts; + BasicBlock &EntryBB = F.getEntryBlock(); + for (Instruction &I : *BB) { + if (ProxiedConsts.size() >= 4) + break; + if (auto *SI = dyn_cast(&I)) { + AllocaInst *PtrAI = getBaseColdAlloca(SI->getPointerOperand()); + auto *C = dyn_cast(SI->getValueOperand()); + if (PtrAI && ReferencedColdAllocas.contains(PtrAI) && C && + !C->getType()->isTokenTy()) { + auto *Proxy = SelectInst::Create( + ConstantInt::getTrue(F.getContext()), C, C, "cold.const", + EntryBB.getTerminator()->getIterator()); + SI->setOperand(0, Proxy); + ProxiedConsts.push_back(Proxy); + } + } + } + + CodeExtractorAnalysisCache CEAC(F); + CodeExtractor CE( + ArrayRef{BB}, /*DT=*/nullptr, + /*AggregateArgs=*/false, /*BFI=*/nullptr, /*BPI=*/nullptr, + /*AC=*/nullptr, /*AllowVarArgs=*/false, /*AllowAlloca=*/true, + /*AllocationBlock=*/nullptr, /*DeallocationBlocks=*/{}, + (OrigII ? "cold.invoke." : "cold.") + std::to_string(++ColdIdx)); + SetVector Inputs, Outputs, Sinks; + CE.findInputsOutputs(Inputs, Outputs, Sinks); + bool Profitable = + Outputs.empty() && Inputs.size() <= 6 && + (NonTrivialInsts >= 2 || OrigCallArgSize > Inputs.size()); + Function *OutF = (Profitable && CE.isEligible()) + ? CE.extractCodeRegion(CEAC) + : nullptr; + for (SelectInst *Proxy : ProxiedConsts) { + Value *OrigC = Proxy->getTrueValue(); + Proxy->replaceAllUsesWith(OrigC); + Proxy->eraseFromParent(); + } + if (OutF && !OutF->user_empty()) { + hoistAllocasToEntry(OutF); + auto *OutCI = cast(*OutF->user_begin()); + OutF->addFnAttr(Attribute::Cold); + OutF->addFnAttr(Attribute::NoInline); + OutF->addFnAttr(Attribute::MinSize); + if (OrigII) { + BasicBlock *CodeRepl = OutCI->getParent(); + Instruction *ReplTerm = CodeRepl->getTerminator(); + SmallVector OutArgs(OutCI->args()); + InvokeInst *NewII = + InvokeInst::Create(OutF, NormalDest, UnwindDest, OutArgs, "", + ReplTerm->getIterator()); + NewII->setCallingConv(OutCI->getCallingConv()); + NewII->setAttributes(OutCI->getAttributes()); + NewII->setDebugLoc(OrigDL); + NewII->addFnAttr(Attribute::Cold); + NewII->addFnAttr(Attribute::NoInline); + OutCI->eraseFromParent(); + ReplTerm->eraseFromParent(); + for (auto &[PN, Val] : SavedPhis) { + PN->addIncoming(Val, CodeRepl); + } + } else { + OutF->addFnAttr(Attribute::NoReturn); + OutCI->addFnAttr(Attribute::Cold); + OutCI->addFnAttr(Attribute::NoInline); + OutCI->addFnAttr(Attribute::NoReturn); + } + Changed = true; + } else { + for (AllocaInst *AI : ColdAllocas) + AI->moveBefore(F.getEntryBlock().getFirstInsertionPt()); + if (OrigII) { + InvokeInst *RestoredII = InvokeInst::Create( + InnerCI->getFunctionType(), InnerCI->getCalledOperand(), + NormalDest, UnwindDest, InnerArgs, Bundles, InnerCI->getName(), + TempUnreach->getIterator()); + RestoredII->setCallingConv(InnerCI->getCallingConv()); + RestoredII->setAttributes(InnerCI->getAttributes()); + RestoredII->setDebugLoc(OrigDL); + InnerCI->eraseFromParent(); + TempUnreach->eraseFromParent(); + for (auto &[PN, Val] : SavedPhis) { + PN->addIncoming(Val, BB); + } + } + } + } + } + return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all(); + } +}; +} // namespace + static codegen::RegisterCodeGenFlags CGF; typedef struct LLVMOpaqueTargetMachine *LLVMTargetMachineRef; @@ -961,6 +1345,23 @@ extern "C" LLVMRustResult LLVMRustOptimize( PB.registerPipelineStartEPCallback(PipelineStartEPCallback); PB.registerOptimizerLastEPCallback(OptimizerLastEPCallback); + static const bool ExpColdOutline = []() { + const char *Env = std::getenv("RUSTC_EXP_COLD_OUTLINE"); + return Env && std::string(Env) == "1"; + }(); + if ((ExpColdOutline || + GEnableColdOutline.load(std::memory_order_relaxed)) && + OptLevel != OptimizationLevel::O0) { + PB.registerPipelineEarlySimplificationEPCallback( + [](ModulePassManager &EarlyMPM, OptimizationLevel Level, + ThinOrFullLTOPhase /*Phase*/) { + if (Level == OptimizationLevel::O0) + return; + EarlyMPM.addPass(ColdInvokeOutlinerPass()); + EarlyMPM.addPass(MergeFunctionsPass()); + }); + } + // The pre-link pipelines don't support O0 and require using // buildO0DefaultPipeline() instead. At the same time, the LTO pipelines // do support O0 and using them is required. diff --git a/compiler/rustc_mir_transform/src/cold_block_outliner.rs b/compiler/rustc_mir_transform/src/cold_block_outliner.rs new file mode 100644 index 0000000000000..4670cf2b07d9a --- /dev/null +++ b/compiler/rustc_mir_transform/src/cold_block_outliner.rs @@ -0,0 +1,455 @@ +//! Experimental MIR pass (`ColdBlockOutliner`) that prepares cold panic, error, and caught-unwind +//! drop regions for outlining and deduplication across generic instantiations. +//! +//! Together with `KnownVariantDrops` and pre-inliner hot/cold splitting, this pass addresses both: +//! - `c8_coldpayload`: caught-panic payload drops (`Box`) on the nonzero path of +//! `catch_unwind`, and +//! - `C10`: cold `panic!` / `assert!` formatting blocks and `?` (`FromResidual::from_residual`) +//! error paths inside generic or `#[inline]` functions. +//! +//! Specifically, `ColdBlockOutliner`: +//! 1. Computes cold basic blocks (`find_cold_blocks`), marking `Unreachable` blocks, calls to +//! `#[cold]` functions, caught-panic paths after `catch_unwind`, and `FromResidual::from_residual` +//! (`?` break) paths, then propagating coldness backwards through the CFG. +//! 2. Sinks pure cold-only setup statements out of hot predecessor blocks into cold entry blocks so +//! they do not bloat the hot path or increase the live-in parameter count of outlined regions. +//! 3. Copies multi-block `Copy` locals into cold-local temporaries before taking shared references +//! in cold blocks so that `non_ssa_locals` keeps the hot local in an SSA register. +//! +//! Enabled via `-Zmir-enable-passes=+ColdBlockOutliner` or `RUSTC_EXP_COLD_OUTLINE=1`. + +use std::sync::LazyLock; + +use rustc_attr_ir::lang_items::LangItem; +use rustc_index::IndexVec; +use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags; +use rustc_middle::mir::visit::{PlaceContext, Visitor}; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, TyCtxt}; +use rustc_span::sym; +use tracing::debug; + +use crate::PassPolicy; + +static COLD_OUTLINE: LazyLock = + LazyLock::new(|| std::env::var("RUSTC_EXP_COLD_OUTLINE").is_ok_and(|v| v == "1")); + +static COLD_CATCH_UNWIND: LazyLock = + LazyLock::new(|| std::env::var("RUSTC_EXP_COLD_CATCH_UNWIND").is_ok_and(|v| v == "1")); + +pub(super) struct ColdBlockOutliner; + +pub(crate) fn is_enabled(tcx: TyCtxt<'_>) -> bool { + if *COLD_OUTLINE { + return true; + } + tcx.sess + .opts + .unstable_opts + .mir_enable_passes + .iter() + .rev() + .find(|(name, _)| name == "ColdBlockOutliner") + .is_some_and(|(_, enabled)| *enabled) +} + +impl<'tcx> crate::MirPass<'tcx> for ColdBlockOutliner { + fn policy(&self, _ctx: &crate::PassCtx<'_>) -> PassPolicy { + PassPolicy::optional(*COLD_OUTLINE) + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let def_id = body.source.def_id(); + if tcx.is_coroutine(def_id) { + return; + } + + let cold_blocks = find_cold_blocks(tcx, body); + if !cold_blocks.iter().any(|&is_cold| is_cold) { + return; + } + + let mut changed = false; + changed |= sink_cold_only_statements(body, &cold_blocks); + changed |= decouple_cold_borrows(tcx, body, &cold_blocks); + + if changed { + debug!(?def_id, "ColdBlockOutliner modified MIR body"); + body.basic_blocks.invalidate_cfg_cache(); + } + } +} + +/// 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: &Body<'tcx>, + cold_blocks: &mut IndexVec, +) { + for bb_data in mir.basic_blocks.iter() { + let 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 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 { + StatementKind::Assign(assign) => assign.0.local == result, + StatementKind::SetDiscriminant { place, .. } => place.local == result, + StatementKind::StorageLive(local) | 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; + } + } +} + +pub(crate) 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)) +} + +/// Identifies cold basic blocks in `mir`, including blocks ending in `Unreachable`, calls to +/// `#[cold]` functions, `FromResidual::from_residual` (`?` residual) calls, and caught-panic +/// arms of `catch_unwind`, then propagates coldness backwards through postorder traversal. +pub(crate) fn find_cold_blocks<'tcx>( + tcx: TyCtxt<'tcx>, + mir: &Body<'tcx>, +) -> IndexVec { + let local_decls = &mir.local_decls; + let mut cold_blocks: IndexVec = IndexVec::from_elem(false, &mir.basic_blocks); + + if *COLD_CATCH_UNWIND || is_enabled(tcx) { + mark_caught_panic_blocks_cold(tcx, mir, &mut cold_blocks); + } + + let mark_residual_cold = is_enabled(tcx); + + for (bb, bb_data) in traversal::postorder(mir) { + if bb_data.is_cleanup { + continue; + } + let terminator = bb_data.terminator(); + + match terminator.kind { + TerminatorKind::Call { ref func, .. } | TerminatorKind::TailCall { ref func, .. } + if let ty::FnDef(def_id, ..) = *func.ty(local_decls, tcx).kind() => + { + let attrs = tcx.codegen_fn_attrs(def_id); + if attrs.flags.contains(CodegenFnAttrFlags::COLD) + || (mark_residual_cold && is_try_from_residual(tcx, def_id)) + { + cold_blocks[bb] = true; + continue; + } + } + + TerminatorKind::Unreachable => { + cold_blocks[bb] = true; + continue; + } + + _ => {} + } + + // Propagate coldness backwards across normal (non-cleanup) edges. + let mut normal_succs = terminator.successors().filter(|&s| !mir.basic_blocks[s].is_cleanup); + if let Some(first) = normal_succs.next() + && cold_blocks[first] + && normal_succs.all(|s| cold_blocks[s]) + { + cold_blocks[bb] = true; + } + } + + cold_blocks +} + +/// Counts uses and definitions of each `Local` across the MIR body. +struct LocalUseCounts { + uses_outside_block: IndexVec, + uses_in_target_block: IndexVec, + target_block: BasicBlock, + source_block: BasicBlock, +} + +impl<'tcx> Visitor<'tcx> for LocalUseCounts { + fn visit_local(&mut self, local: Local, context: PlaceContext, location: Location) { + if context.is_storage_marker() { + return; + } + if location.block == self.target_block { + self.uses_in_target_block[local] += 1; + } else if location.block != self.source_block { + self.uses_outside_block[local] += 1; + } + } +} + +fn is_pure_sinkable_rvalue<'tcx>(rvalue: &Rvalue<'tcx>) -> bool { + matches!( + rvalue, + Rvalue::Use(..) + | Rvalue::Ref(_, BorrowKind::Shared, _) + | Rvalue::RawPtr(_, _) + | Rvalue::Cast(_, _, _) + | Rvalue::BinaryOp(_, _) + | Rvalue::UnaryOp(_, _) + | Rvalue::Discriminant(_) + | Rvalue::Aggregate(_, _) + ) +} + +fn rvalue_reads_local<'tcx>(rvalue: &Rvalue<'tcx>, target_local: Local) -> bool { + struct ReadsLocal { + target_local: Local, + found: bool, + } + impl<'tcx> Visitor<'tcx> for ReadsLocal { + fn visit_local(&mut self, local: Local, _context: PlaceContext, _location: Location) { + if local == self.target_local { + self.found = true; + } + } + } + let mut finder = ReadsLocal { target_local, found: false }; + finder.visit_rvalue(rvalue, START_BLOCK.start_location()); + finder.found +} + +/// Sinks pure temporary assignments out of hot blocks into a cold successor block when the +/// assigned local is only used inside that cold successor block and the cold successor has a +/// single predecessor. +fn sink_cold_only_statements<'tcx>( + body: &mut Body<'tcx>, + cold_blocks: &IndexVec, +) -> bool { + let mut pred_counts: IndexVec = IndexVec::from_elem(0, &body.basic_blocks); + for data in body.basic_blocks.iter() { + for succ in data.terminator().successors() { + pred_counts[succ] += 1; + } + } + + let mut candidates = Vec::new(); + for (bb, data) in body.basic_blocks.iter_enumerated() { + if data.is_cleanup || cold_blocks[bb] || data.statements.is_empty() { + continue; + } + let TerminatorKind::SwitchInt { ref discr, ref targets } = data.terminator().kind else { + continue; + }; + let discr_local = discr.place().and_then(|p| p.as_local()); + let mut cold_targets = + targets.all_targets().iter().copied().filter(|&t| { + cold_blocks[t] && !body.basic_blocks[t].is_cleanup && pred_counts[t] == 1 + }); + let Some(cold_target) = cold_targets.next() else { + continue; + }; + if cold_targets.next().is_some() { + continue; + } + candidates.push((bb, cold_target, discr_local)); + } + + let mut changed = false; + for (bb_hot, bb_cold, discr_local) in candidates { + let mut counts = LocalUseCounts { + uses_outside_block: IndexVec::from_elem(0, &body.local_decls), + uses_in_target_block: IndexVec::from_elem(0, &body.local_decls), + target_block: bb_cold, + source_block: bb_hot, + }; + counts.visit_body(body); + + let statements = &body.basic_blocks[bb_hot].statements; + let mut to_sink_indices = Vec::new(); + for (idx, stmt) in statements.iter().enumerate() { + let StatementKind::Assign(ref assign) = stmt.kind else { + continue; + }; + let (place, ref rvalue) = **assign; + let Some(local) = place.as_local() else { + continue; + }; + if Some(local) == discr_local || local == RETURN_PLACE { + continue; + } + if !is_pure_sinkable_rvalue(rvalue) { + continue; + } + if counts.uses_outside_block[local] != 0 || counts.uses_in_target_block[local] == 0 { + continue; + } + // Ensure no later statement in `bb_hot` reads `local` or writes any local read by `rvalue`. + let used_or_clobbered_later = + statements[idx + 1..].iter().any(|later| match &later.kind { + StatementKind::Assign(later_assign) => { + let (later_place, ref later_rvalue) = **later_assign; + rvalue_reads_local(later_rvalue, local) + || later_place.local == local + || rvalue_reads_local(rvalue, later_place.local) + } + StatementKind::SetDiscriminant { place: later_place, .. } => { + later_place.local == local || rvalue_reads_local(rvalue, later_place.local) + } + StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => { + *l == local || rvalue_reads_local(rvalue, *l) + } + _ => true, + }); + if !used_or_clobbered_later { + to_sink_indices.push(idx); + } + } + + if to_sink_indices.is_empty() { + continue; + } + + let bbs = body.basic_blocks.as_mut(); + let mut sunken = Vec::with_capacity(to_sink_indices.len()); + for &idx in to_sink_indices.iter().rev() { + sunken.push(bbs[bb_hot].statements.remove(idx)); + } + sunken.reverse(); + sunken.append(&mut bbs[bb_cold].statements); + bbs[bb_cold].statements = sunken; + changed = true; + } + + changed +} + +/// Copies `Copy` locals that are live across non-cold blocks into fresh cold-local temporaries +/// before taking a shared reference (`&local`) inside a cold block (such as `panic!("{x}")`). +/// +/// Without this transformation, `rustc_codegen_ssa::mir::analyze::non_ssa_locals` sees the shared +/// borrow of `x` in the cold block and forces `x` into a stack `alloca` in the function entry +/// block (`start`), emitting `alloca` + `store` on the hot path and preventing `CodeExtractor` +/// from sinking the `alloca` into the outlined `.cold` function. +fn decouple_cold_borrows<'tcx>( + tcx: TyCtxt<'tcx>, + body: &mut Body<'tcx>, + cold_blocks: &IndexVec, +) -> bool { + let num_locals = body.local_decls.len(); + let mut local_blocks: IndexVec> = + IndexVec::from_elem_n(None, num_locals); + let mut multi_block_locals = rustc_index::bit_set::DenseBitSet::new_empty(num_locals); + + for arg in 1..=body.arg_count { + multi_block_locals.insert(Local::from_usize(arg)); + } + multi_block_locals.insert(RETURN_PLACE); + + struct LocalBlockTracker<'a> { + local_blocks: &'a mut IndexVec>, + multi_block_locals: &'a mut rustc_index::bit_set::DenseBitSet, + } + + impl<'tcx> Visitor<'tcx> for LocalBlockTracker<'_> { + fn visit_local(&mut self, local: Local, _context: PlaceContext, location: Location) { + if self.multi_block_locals.contains(local) { + return; + } + match self.local_blocks[local] { + None => self.local_blocks[local] = Some(location.block), + Some(prev_bb) if prev_bb != location.block => { + self.multi_block_locals.insert(local); + } + Some(_) => {} + } + } + } + + LocalBlockTracker { + local_blocks: &mut local_blocks, + multi_block_locals: &mut multi_block_locals, + } + .visit_body(body); + + let typing_env = body.typing_env(tcx); + let mut changed = false; + + for bb in body.basic_blocks.indices() { + if !cold_blocks[bb] || body.basic_blocks[bb].is_cleanup { + continue; + } + let stmt_len = body.basic_blocks[bb].statements.len(); + if stmt_len == 0 { + continue; + } + + let mut new_stmts = Vec::with_capacity(stmt_len); + let old_stmts = std::mem::take(&mut body.basic_blocks.as_mut()[bb].statements); + + for mut stmt in old_stmts { + let mut copy_prefix: Option<(Local, Local, SourceInfo)> = None; + if let StatementKind::Assign(ref mut assign) = stmt.kind { + let (_, ref mut rvalue) = **assign; + if let Rvalue::Ref(region, BorrowKind::Shared, ref mut borrowed_place) = *rvalue + && let Some(local) = borrowed_place.as_local() + && local.as_usize() < num_locals + && multi_block_locals.contains(local) + { + let ty = body.local_decls[local].ty; + if tcx.type_is_copy_modulo_regions(typing_env, ty) { + let source_info = stmt.source_info; + let cold_local = + body.local_decls.push(LocalDecl::with_source_info(ty, source_info)); + *borrowed_place = Place::from(cold_local); + let _ = region; + copy_prefix = Some((cold_local, local, source_info)); + } + } + } + if let Some((cold_local, orig_local, source_info)) = copy_prefix { + new_stmts.push(Statement::new(source_info, StatementKind::StorageLive(cold_local))); + new_stmts.push(Statement::new( + source_info, + StatementKind::Assign(Box::new(( + Place::from(cold_local), + Rvalue::Use(Operand::Copy(Place::from(orig_local)), WithRetag::No), + ))), + )); + changed = true; + } + new_stmts.push(stmt); + } + + body.basic_blocks.as_mut()[bb].statements = new_stmts; + } + + changed +} diff --git a/compiler/rustc_mir_transform/src/inline.rs b/compiler/rustc_mir_transform/src/inline.rs index 367e16daa8232..a3bc21d386acb 100644 --- a/compiler/rustc_mir_transform/src/inline.rs +++ b/compiler/rustc_mir_transform/src/inline.rs @@ -8,8 +8,8 @@ use rustc_attr_ir::lang_items::LangItem; use rustc_attr_ir::{InlineAttr, OptimizeAttr}; use rustc_hir::def::DefKind; use rustc_hir::def_id::DefId; -use rustc_index::Idx; use rustc_index::bit_set::DenseBitSet; +use rustc_index::{Idx, IndexVec}; use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs; use rustc_middle::mir::visit::*; use rustc_middle::mir::*; @@ -499,6 +499,7 @@ fn process_blocks<'tcx, I: Inliner<'tcx>>( caller_body: &mut Body<'tcx>, blocks: Range, ) { + let mut caller_cold: Option> = None; for bb in blocks { let bb_data = &caller_body[bb]; if bb_data.is_cleanup { @@ -512,13 +513,14 @@ fn process_blocks<'tcx, I: Inliner<'tcx>>( let span = trace_span!("process_blocks", %callsite.callee, ?bb); let _guard = span.enter(); - match try_inlining(inliner, caller_body, &callsite) { + match try_inlining(inliner, caller_body, &callsite, &mut caller_cold) { Err(reason) => { debug!("not-inlined {} [{}]", callsite.callee, reason); inliner.on_inline_failure(&callsite, reason); } Ok(new_blocks) => { debug!("inlined {}", callsite.callee); + caller_cold = None; inliner.on_inline_success(&callsite, caller_body, new_blocks); } } @@ -601,6 +603,7 @@ fn try_inlining<'tcx, I: Inliner<'tcx>>( inliner: &I, caller_body: &mut Body<'tcx>, callsite: &CallSite<'tcx>, + caller_cold: &mut Option>, ) -> Result, &'static str> { let tcx = inliner.tcx(); check_mir_is_available(inliner, caller_body, callsite.callee)?; @@ -610,6 +613,18 @@ fn try_inlining<'tcx, I: Inliner<'tcx>>( check_inline::is_inline_valid_on_fn(tcx, callsite.callee.def_id())?; check_codegen_attributes(inliner, callsite, callee_attrs)?; + if crate::cold_block_outliner::is_enabled(tcx) + && !matches!(callee_attrs.inline, InlineAttr::Always | InlineAttr::Force { .. }) + && (matches!(callsite.callee.def, InstanceKind::Shim(ty::ShimKind::DropGlue(..))) + || crate::cold_block_outliner::is_try_from_residual(tcx, callsite.callee.def_id())) + { + let cold_blocks = caller_cold + .get_or_insert_with(|| crate::cold_block_outliner::find_cold_blocks(tcx, caller_body)); + if cold_blocks[callsite.block] { + return Err("callsite is in a cold block"); + } + } + let terminator = caller_body[callsite.block].terminator.as_ref().unwrap(); let TerminatorKind::Call { args, destination, .. } = &terminator.kind else { bug!() }; let destination_ty = destination.ty(&caller_body.local_decls, tcx).ty; diff --git a/compiler/rustc_mir_transform/src/known_variant_drops.rs b/compiler/rustc_mir_transform/src/known_variant_drops.rs new file mode 100644 index 0000000000000..fee3f4d32c846 --- /dev/null +++ b/compiler/rustc_mir_transform/src/known_variant_drops.rs @@ -0,0 +1,598 @@ +//! Experimental, local prototype: specializes `Drop` terminators of enum locals on the incoming +//! edges where the variant of the local is known. +//! +//! After `catch_unwind` is inlined, its `Result>` is built in an `Ok` and +//! an `Err` block, joined, matched, and then dropped as a whole. The whole-value drop takes the +//! address of the local. That keeps the local out of `JumpThreading` (which ignores locals whose +//! address may be observed) and out of SROA in LLVM, so the drop glue of the `Err` payload sits on +//! the join path, looks hot, and gets inlined into every caller. +//! +//! On an edge where the variant is known, the drop is replaced by a drop of the only field of that +//! variant that needs dropping (moved to a fresh local first), or by a `goto` if no field needs +//! dropping. When every such drop is rewritten, the local is no longer address-taken and +//! `JumpThreading` can thread the match through the join, leaving the payload drop only on the +//! path that caught a panic. +//! +//! Only runs with `-Zmir-enable-passes=+KnownVariantDrops`. + +use rustc_abi::{FieldIdx, VariantIdx}; +use rustc_const_eval::const_eval::DummyMachine; +use rustc_const_eval::interpret::InterpCx; +use rustc_data_structures::fx::FxIndexMap; +use rustc_index::bit_set::DenseBitSet; +use rustc_index::{IndexSlice, IndexVec}; +use rustc_middle::mir::visit::{MutatingUseContext, PlaceContext, Visitor}; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, Ty, TyCtxt}; +use rustc_span::{DUMMY_SP, bug}; +use tracing::{debug, trace}; + +use crate::PassPolicy; +use crate::simplify::simplify_cfg; + +pub(super) struct KnownVariantDrops; + +/// Drop blocks are duplicated once per known variant; larger blocks are left alone. +const MAX_DROP_BLOCK_STATEMENTS: usize = 32; + +impl<'tcx> crate::MirPass<'tcx> for KnownVariantDrops { + fn policy(&self, _ctx: &crate::PassCtx<'_>) -> PassPolicy { + // Experimental: off unless requested with `-Zmir-enable-passes`. + PassPolicy::optional(false) + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let def_id = body.source.def_id(); + // Same restriction as `JumpThreading`, which this pass feeds. + if tcx.is_coroutine(def_id) { + return; + } + + let typing_env = body.typing_env(tcx); + let candidates = collect_candidates(tcx, typing_env, body); + if candidates.is_empty() { + return; + } + + let rewrites = { + let analysis = Analysis::new(typing_env, tcx, body, &candidates); + let entry_states = analysis.entry_states(); + analysis.find_rewrites(&entry_states) + }; + if rewrites.is_empty() { + return; + } + + debug!(?def_id, rewrites = rewrites.len(), "specializing drops of known variants"); + apply_rewrites(tcx, body, &candidates, &rewrites); + simplify_cfg(tcx, body); + } +} + +/// What dropping the whole enum does when it holds a given variant. +#[derive(Copy, Clone, Debug)] +enum DropAction<'tcx> { + /// No field of the variant needs dropping. + Nothing, + /// Exactly one field of the variant needs dropping. + Field(FieldIdx, Ty<'tcx>), + /// Several fields need dropping; the whole-value drop is kept. + Whole, +} + +struct Candidate<'tcx> { + local: Local, + adt: ty::AdtDef<'tcx>, + actions: IndexVec>, + /// Discriminant bit pattern of each variant, as compared by `SwitchInt` on + /// `discriminant(local)`. + discrs: Vec<(u128, VariantIdx)>, +} + +/// Enum locals without a destructor that are dropped whole outside cleanup blocks and whose +/// address is not observed otherwise. +fn collect_candidates<'tcx>( + tcx: TyCtxt<'tcx>, + typing_env: ty::TypingEnv<'tcx>, + body: &Body<'tcx>, +) -> Vec> { + let mut dropped = DenseBitSet::new_empty(body.local_decls.len()); + for data in body.basic_blocks.iter() { + if data.is_cleanup { + continue; + } + if let TerminatorKind::Drop { place, drop: None, .. } = data.terminator().kind + && let Some(local) = place.as_local() + { + dropped.insert(local); + } + } + if dropped.is_empty() { + return Vec::new(); + } + + let mut observed = AddressObserved { locals: DenseBitSet::new_empty(body.local_decls.len()) }; + observed.visit_body(body); + + let mut candidates = Vec::new(); + for local in dropped.iter() { + if observed.locals.contains(local) { + continue; + } + let ty::Adt(adt, args) = *body.local_decls[local].ty.kind() else { continue }; + if !adt.is_enum() || adt.has_dtor(tcx) { + continue; + } + let actions: IndexVec> = adt + .variants() + .iter() + .map(|variant| { + let mut action = DropAction::Nothing; + for (field_idx, field) in variant.fields.iter_enumerated() { + let field_ty = field.ty(tcx, args); + let field_ty = tcx + .try_normalize_erasing_regions(typing_env, field_ty) + .unwrap_or(field_ty.skip_norm_wip()); + if !field_ty.needs_drop(tcx, typing_env) { + continue; + } + action = match action { + DropAction::Nothing => DropAction::Field(field_idx, field_ty), + DropAction::Field(..) | DropAction::Whole => DropAction::Whole, + }; + } + action + }) + .collect(); + if actions.iter().all(|action| matches!(action, DropAction::Whole)) { + continue; + } + let discrs = adt.discriminants(tcx).map(|(variant, discr)| (discr.val, variant)).collect(); + trace!(?local, ?actions, "candidate"); + candidates.push(Candidate { local, adt, actions, discrs }); + } + candidates +} + +/// Locals whose address may be observed by anything other than a drop of the whole local. +struct AddressObserved { + locals: DenseBitSet, +} + +impl<'tcx> Visitor<'tcx> for AddressObserved { + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _location: Location) { + let whole_drop = context == PlaceContext::MutatingUse(MutatingUseContext::Drop) + && place.projection.is_empty(); + if context.may_observe_address() && !place.is_indirect() && !whole_drop { + self.locals.insert(place.local); + } + } +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +enum Val { + /// No execution reaches this point. + Unreached, + Known(VariantIdx), + Unknown, +} + +impl Val { + fn join(self, other: Val) -> Val { + match (self, other) { + (Val::Unreached, v) | (v, Val::Unreached) => v, + (Val::Known(a), Val::Known(b)) if a == b => self, + _ => Val::Unknown, + } + } + + /// Narrows `self` with what taking a switch edge implies. + fn meet(self, other: Val) -> Val { + match (self, other) { + (Val::Unknown, v) | (v, Val::Unknown) => v, + (Val::Known(a), Val::Known(b)) if a == b => self, + _ => Val::Unreached, + } + } +} + +#[derive(Copy, Clone, Debug)] +struct Edge { + from: BasicBlock, + /// Index into `SwitchTargets::all_targets` when `from` ends in a `SwitchInt`. + slot: Option, + target: BasicBlock, +} + +struct Rewrite { + drop_block: BasicBlock, + candidate: usize, + variant: VariantIdx, + edges: Vec, +} + +/// Forgets the variant of every candidate that the visited code may write. +struct Clobber<'a> { + index: &'a IndexSlice>, + state: &'a mut [Val], +} + +impl<'tcx> Visitor<'tcx> for Clobber<'_> { + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _location: Location) { + if context.is_mutating_use() + && let Some(i) = self.index[place.local] + { + self.state[i] = Val::Unknown; + } + } + + fn visit_local(&mut self, local: Local, context: PlaceContext, _location: Location) { + if (context.is_mutating_use() || context.is_storage_marker()) + && let Some(i) = self.index[local] + { + self.state[i] = Val::Unknown; + } + } +} + +/// Whether `statement` may change the value or the storage of `local`. +fn writes_local<'tcx>(statement: &Statement<'tcx>, local: Local) -> bool { + struct Finder { + local: Local, + found: bool, + } + + impl<'tcx> Visitor<'tcx> for Finder { + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _location: Location) { + if place.local == self.local && context.is_mutating_use() { + self.found = true; + } + } + + fn visit_local(&mut self, local: Local, context: PlaceContext, _location: Location) { + if local == self.local && (context.is_mutating_use() || context.is_storage_marker()) { + self.found = true; + } + } + } + + let mut finder = Finder { local, found: false }; + finder.visit_statement(statement, START_BLOCK.start_location()); + finder.found +} + +struct Analysis<'a, 'tcx> { + body: &'a Body<'tcx>, + ecx: InterpCx<'tcx, DummyMachine>, + candidates: &'a [Candidate<'tcx>], + /// Candidate index of each local. + index: IndexVec>, +} + +impl<'a, 'tcx> Analysis<'a, 'tcx> { + fn new( + typing_env: ty::TypingEnv<'tcx>, + tcx: TyCtxt<'tcx>, + body: &'a Body<'tcx>, + candidates: &'a [Candidate<'tcx>], + ) -> Self { + let mut index = IndexVec::from_elem(None, &body.local_decls); + for (i, candidate) in candidates.iter().enumerate() { + index[candidate.local] = Some(i); + } + Analysis { + body, + ecx: InterpCx::new(tcx, DUMMY_SP, typing_env, DummyMachine), + candidates, + index, + } + } + + fn assigned_variant( + &self, + candidate: &Candidate<'tcx>, + rvalue: &Rvalue<'tcx>, + ) -> Option { + match rvalue { + Rvalue::Aggregate(kind, _) => match **kind { + AggregateKind::Adt(did, variant, _, _, None) if did == candidate.adt.did() => { + Some(variant) + } + _ => None, + }, + // `Ok(())` and other field-less values are constants after GVN. + Rvalue::Use(Operand::Constant(constant), _) => { + let op = self + .ecx + .eval_mir_constant(&constant.const_, constant.span, None) + .discard_err()?; + self.ecx.read_discriminant(&op).discard_err() + } + _ => None, + } + } + + fn apply_statement(&self, statement: &Statement<'tcx>, location: Location, state: &mut [Val]) { + Clobber { index: &self.index, state: &mut *state }.visit_statement(statement, location); + match &statement.kind { + StatementKind::Assign(assign) => { + let (place, rvalue) = &**assign; + if let Some(local) = place.as_local() + && let Some(i) = self.index[local] + { + state[i] = self + .assigned_variant(&self.candidates[i], rvalue) + .map_or(Val::Unknown, Val::Known); + } + } + StatementKind::SetDiscriminant { place, variant_index } => { + if let Some(local) = place.as_local() + && let Some(i) = self.index[local] + { + state[i] = Val::Known(*variant_index); + } + } + _ => {} + } + } + + fn apply_statements(&self, bb: BasicBlock, state: &mut [Val]) { + for (statement_index, statement) in self.body.basic_blocks[bb].statements.iter().enumerate() + { + self.apply_statement(statement, Location { block: bb, statement_index }, state); + } + } + + /// For `switchInt(_d)` where the last write of `_d` in `bb` is `_d = discriminant(P)` and + /// neither `_d` nor `P` is written after it, returns P's candidate index and the value of P + /// implied by each slot of `targets.all_targets()`. + fn switch_refinement( + &self, + bb: BasicBlock, + discr: &Operand<'tcx>, + targets: &SwitchTargets, + ) -> Option<(usize, Vec)> { + let discr_local = discr.place()?.as_local()?; + let statements = &self.body.basic_blocks[bb].statements; + let (position, i) = + statements.iter().enumerate().rev().find_map(|(position, statement)| { + if !writes_local(statement, discr_local) { + return None; + } + let found = if let StatementKind::Assign(assign) = &statement.kind + && let (place, Rvalue::Discriminant(enum_place)) = &**assign + && place.as_local() == Some(discr_local) + && let Some(enum_local) = enum_place.as_local() + && let Some(i) = self.index[enum_local] + { + Some((position, i)) + } else { + None + }; + Some(found) + })??; + let candidate = &self.candidates[i]; + if statements[position + 1..] + .iter() + .any(|s| writes_local(s, discr_local) || writes_local(s, candidate.local)) + { + return None; + } + + let mut remaining: Vec = candidate.discrs.iter().map(|&(_, v)| v).collect(); + let mut values = Vec::with_capacity(targets.all_targets().len()); + for (value, _) in targets.iter() { + // Give up on values that match no variant instead of guessing what they mean. + let &(_, variant) = candidate.discrs.iter().find(|&&(d, _)| d == value)?; + values.push(Val::Known(variant)); + remaining.retain(|&v| v != variant); + } + values.push(match remaining[..] { + [] => Val::Unreached, + [variant] => Val::Known(variant), + _ => Val::Unknown, + }); + Some((i, values)) + } + + /// Calls `f` with the state on every edge from `bb` into a non-cleanup block, given the + /// state before `bb`'s terminator. + fn for_each_edge(&self, bb: BasicBlock, state: &[Val], mut f: impl FnMut(Edge, &[Val])) { + let data = &self.body.basic_blocks[bb]; + let terminator = data.terminator(); + let location = Location { block: bb, statement_index: data.statements.len() }; + let mut exit = state.to_vec(); + Clobber { index: &self.index, state: &mut exit }.visit_terminator(terminator, location); + + if let TerminatorKind::SwitchInt { discr, targets } = &terminator.kind { + let refinement = self.switch_refinement(bb, discr, targets); + let mut edge_state = exit.clone(); + for (slot, &target) in targets.all_targets().iter().enumerate() { + if self.body.basic_blocks[target].is_cleanup { + continue; + } + edge_state.copy_from_slice(&exit); + if let Some((i, values)) = &refinement { + edge_state[*i] = exit[*i].meet(values[slot]); + } + f(Edge { from: bb, slot: Some(slot), target }, &edge_state); + } + } else { + for target in terminator.successors() { + if self.body.basic_blocks[target].is_cleanup { + continue; + } + f(Edge { from: bb, slot: None, target }, &exit); + } + } + } + + /// Forward dataflow to a fixpoint. `None` means that no edge reaches the block. + fn entry_states(&self) -> IndexVec>> { + let mut entry: IndexVec>> = + IndexVec::from_elem(None, &self.body.basic_blocks); + entry[START_BLOCK] = Some(vec![Val::Unknown; self.candidates.len()]); + let rpo = self.body.basic_blocks.reverse_postorder(); + let mut changed = true; + while changed { + changed = false; + for &bb in rpo { + if self.body.basic_blocks[bb].is_cleanup { + continue; + } + let Some(mut state) = entry[bb].clone() else { continue }; + self.apply_statements(bb, &mut state); + self.for_each_edge(bb, &state, |edge, edge_state| { + let target_state = &mut entry[edge.target]; + match target_state { + Some(old) => { + for (old, &new) in old.iter_mut().zip(edge_state) { + let joined = old.join(new); + changed |= joined != *old; + *old = joined; + } + } + None => { + *target_state = Some(edge_state.to_vec()); + changed = true; + } + } + }); + } + } + entry + } + + fn find_rewrites(&self, entry: &IndexSlice>>) -> Vec { + // Drop block -> (candidate index, incoming edges with their states). + let mut drop_blocks: FxIndexMap)>)> = + FxIndexMap::default(); + for (bb, data) in self.body.basic_blocks.iter_enumerated() { + if data.is_cleanup || data.statements.len() > MAX_DROP_BLOCK_STATEMENTS { + continue; + } + if let TerminatorKind::Drop { place, replace: false, drop: None, .. } = + data.terminator().kind + && let Some(local) = place.as_local() + && let Some(i) = self.index[local] + { + drop_blocks.insert(bb, (i, Vec::new())); + } + } + if drop_blocks.is_empty() { + return Vec::new(); + } + + for &bb in self.body.basic_blocks.reverse_postorder() { + if self.body.basic_blocks[bb].is_cleanup { + continue; + } + let Some(state) = &entry[bb] else { continue }; + let mut state = state.clone(); + self.apply_statements(bb, &mut state); + self.for_each_edge(bb, &state, |edge, edge_state| { + if let Some((_, incoming)) = drop_blocks.get_mut(&edge.target) { + incoming.push((edge, edge_state.to_vec())); + } + }); + } + + let mut rewrites = Vec::new(); + for (drop_block, (i, incoming)) in drop_blocks { + let candidate = &self.candidates[i]; + let mut groups: FxIndexMap> = FxIndexMap::default(); + for (edge, mut state) in incoming { + if edge.from == drop_block { + continue; + } + self.apply_statements(drop_block, &mut state); + if let Val::Known(variant) = state[i] + && !matches!(candidate.actions[variant], DropAction::Whole) + { + groups.entry(variant).or_default().push(edge); + } + } + for (variant, edges) in groups { + trace!(?drop_block, ?variant, ?edges, "rewrite"); + rewrites.push(Rewrite { drop_block, candidate: i, variant, edges }); + } + } + rewrites + } +} + +fn apply_rewrites<'tcx>( + tcx: TyCtxt<'tcx>, + body: &mut Body<'tcx>, + candidates: &[Candidate<'tcx>], + rewrites: &[Rewrite], +) { + // Clone every drop block before retargeting any edge, so that each copy keeps the original + // successors. + let mut new_blocks = Vec::with_capacity(rewrites.len()); + for rewrite in rewrites { + let candidate = &candidates[rewrite.candidate]; + let mut data = body.basic_blocks[rewrite.drop_block].clone(); + let source_info = data.terminator().source_info; + let TerminatorKind::Drop { place, target, unwind, replace, drop: None } = + data.terminator().kind + else { + bug!("expected a drop terminator in {:?}", rewrite.drop_block); + }; + let kind = match candidate.actions[rewrite.variant] { + DropAction::Nothing => { + TerminatorKind::Goto { target, loop_hint_attrs: Default::default() } + } + DropAction::Field(field, field_ty) => { + // Dropping a moved-out copy is fine: the local's address is never observed, so + // nothing can rely on the field being dropped in place. + let temp = body.local_decls.push(LocalDecl::new(field_ty, source_info.span)); + let field_place = tcx.mk_place_field( + tcx.mk_place_downcast(place, candidate.adt, rewrite.variant), + field, + field_ty, + ); + let mut statement = Statement::new( + source_info, + StatementKind::Assign(Box::new(( + Place::from(temp), + Rvalue::Use(Operand::Move(field_place), WithRetag::Yes), + ))), + ); + statement.debuginfos.prepend(&mut data.after_last_stmt_debuginfos); + data.statements.push(statement); + TerminatorKind::Drop { + place: Place::from(temp), + target, + unwind, + replace, + drop: None, + } + } + DropAction::Whole => bug!("whole-value drops are not rewritten"), + }; + data.terminator_mut().kind = kind; + new_blocks.push(body.basic_blocks.as_mut().push(data)); + } + + for (rewrite, &new_block) in rewrites.iter().zip(&new_blocks) { + for edge in &rewrite.edges { + let terminator = body.basic_blocks.as_mut()[edge.from].terminator_mut(); + if let Some(slot) = edge.slot { + let TerminatorKind::SwitchInt { targets, .. } = &mut terminator.kind else { + bug!("switch edge from a non-switch terminator in {:?}", edge.from); + }; + let target = &mut targets.all_targets_mut()[slot]; + assert_eq!(*target, rewrite.drop_block); + *target = new_block; + } else { + let drop_block = rewrite.drop_block; + terminator.successors_mut(|successor| { + if *successor == drop_block { + *successor = new_block; + } + }); + } + } + } +} diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs index d73d9a82f1a22..2a6337d60e3ae 100644 --- a/compiler/rustc_mir_transform/src/lib.rs +++ b/compiler/rustc_mir_transform/src/lib.rs @@ -133,6 +133,7 @@ declare_passes! { mod check_null : CheckNull; mod check_packed_ref : CheckPackedRef; mod check_mut_restriction : CheckMutRestriction; + mod cold_block_outliner : ColdBlockOutliner; // This pass is public to allow external drivers to perform MIR cleanup pub mod cleanup_post_borrowck : CleanupPostBorrowck; @@ -160,6 +161,7 @@ declare_passes! { mod instsimplify : InstSimplify { BeforeInline, AfterSimplifyCfg }; mod jump_threading : JumpThreading; mod known_panics_lint : KnownPanicsLint; + mod known_variant_drops : KnownVariantDrops; mod lint_and_remove_uninhabited : LintAndRemoveUninhabited; mod lower_intrinsics : LowerIntrinsics; mod lower_slice_len : LowerSliceLenCalls; @@ -759,6 +761,7 @@ pub(crate) fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<' &dataflow_const_prop::DataflowConstProp, &single_use_consts::SingleUseConsts, &o1(simplify_branches::SimplifyConstCondition::AfterConstProp), + &known_variant_drops::KnownVariantDrops, &jump_threading::JumpThreading, &early_otherwise_branch::EarlyOtherwiseBranch, &simplify_comparison_integral::SimplifyComparisonIntegral, @@ -771,6 +774,7 @@ pub(crate) fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<' &dead_store_elimination::DeadStoreElimination::Final, &dest_prop::DestinationPropagation, &tail_copy_to_move::TailCopyToMove, + &cold_block_outliner::ColdBlockOutliner, &move_elimination::MoveElimination, &simplify::SimplifyLocals::Final, &multiple_return_terminators::MultipleReturnTerminators,