From 5fa7d14bba9dc39a4c1832faf846992a1452578a Mon Sep 17 00:00:00 2001 From: Ben Kimock Date: Sun, 31 Mar 2024 12:31:48 -0400 Subject: [PATCH] Only collect mono items from reachable blocks --- Cargo.lock | 1 + compiler/rustc_codegen_ssa/src/mir/mod.rs | 9 ++- compiler/rustc_interface/src/passes.rs | 1 + compiler/rustc_middle/src/mir/mod.rs | 52 -------------- compiler/rustc_middle/src/mir/traversal.rs | 73 ++++++++++++++++++++ compiler/rustc_middle/src/query/keys.rs | 17 +++++ compiler/rustc_middle/src/query/mod.rs | 10 ++- compiler/rustc_monomorphize/Cargo.toml | 1 + compiler/rustc_monomorphize/src/collector.rs | 17 +++++ 9 files changed, 122 insertions(+), 59 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6168a0a174719..ab903afd9d014 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4374,6 +4374,7 @@ dependencies = [ "rustc_errors", "rustc_fluent_macro", "rustc_hir", + "rustc_index", "rustc_macros", "rustc_middle", "rustc_session", diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs index 387a5366b209b..8a5ad35eb8fa3 100644 --- a/compiler/rustc_codegen_ssa/src/mir/mod.rs +++ b/compiler/rustc_codegen_ssa/src/mir/mod.rs @@ -257,20 +257,19 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>( // Apply debuginfo to the newly allocated locals. fx.debug_introduce_locals(&mut start_bx); - let reachable_blocks = mir.reachable_blocks_in_mono(cx.tcx(), instance); - // The builders will be created separately for each basic block at `codegen_block`. // So drop the builder of `start_llbb` to avoid having two at the same time. drop(start_bx); + let reachable_blocks = cx.tcx().reachable_blocks(instance); + // Codegen the body of each block using reverse postorder for (bb, _) in traversal::reverse_postorder(mir) { if reachable_blocks.contains(bb) { fx.codegen_block(bb); } else { - // This may have references to things we didn't monomorphize, so we - // don't actually codegen the body. We still create the block so - // terminators in other blocks can reference it without worry. + // We want to skip this block, because it's not reachable. But we still create + // the block so terminators in other blocks can reference it. fx.codegen_block_as_unreachable(bb); } } diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs index d763a12f816b0..e6653ea7e4ed7 100644 --- a/compiler/rustc_interface/src/passes.rs +++ b/compiler/rustc_interface/src/passes.rs @@ -631,6 +631,7 @@ pub static DEFAULT_QUERY_PROVIDERS: LazyLock = LazyLock::new(|| { rustc_lint::provide(providers); rustc_symbol_mangling::provide(providers); rustc_codegen_ssa::provide(providers); + rustc_middle::mir::traversal::provide(providers); *providers }); diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs index e5a650c5ac4ae..5862481232e4b 100644 --- a/compiler/rustc_middle/src/mir/mod.rs +++ b/compiler/rustc_middle/src/mir/mod.rs @@ -30,7 +30,6 @@ pub use rustc_ast::Mutability; use rustc_data_structures::fx::FxHashMap; use rustc_data_structures::fx::FxHashSet; use rustc_data_structures::graph::dominators::Dominators; -use rustc_data_structures::stack::ensure_sufficient_stack; use rustc_index::bit_set::BitSet; use rustc_index::{Idx, IndexSlice, IndexVec}; use rustc_serialize::{Decodable, Encodable}; @@ -687,57 +686,6 @@ impl<'tcx> Body<'tcx> { self.injection_phase.is_some() } - /// Finds which basic blocks are actually reachable for a specific - /// monomorphization of this body. - /// - /// This is allowed to have false positives; just because this says a block - /// is reachable doesn't mean that's necessarily true. It's thus always - /// legal for this to return a filled set. - /// - /// Regardless, the [`BitSet::domain_size`] of the returned set will always - /// exactly match the number of blocks in the body so that `contains` - /// checks can be done without worrying about panicking. - /// - /// This is mostly useful because it lets us skip lowering the `false` side - /// of `if ::CONST`, as well as `intrinsics::debug_assertions`. - pub fn reachable_blocks_in_mono( - &self, - tcx: TyCtxt<'tcx>, - instance: Instance<'tcx>, - ) -> BitSet { - let mut set = BitSet::new_empty(self.basic_blocks.len()); - self.reachable_blocks_in_mono_from(tcx, instance, &mut set, START_BLOCK); - set - } - - fn reachable_blocks_in_mono_from( - &self, - tcx: TyCtxt<'tcx>, - instance: Instance<'tcx>, - set: &mut BitSet, - bb: BasicBlock, - ) { - if !set.insert(bb) { - return; - } - - let data = &self.basic_blocks[bb]; - - if let Some((bits, targets)) = Self::try_const_mono_switchint(tcx, instance, data) { - let target = targets.target_for_value(bits); - ensure_sufficient_stack(|| { - self.reachable_blocks_in_mono_from(tcx, instance, set, target) - }); - return; - } - - for target in data.terminator().successors() { - ensure_sufficient_stack(|| { - self.reachable_blocks_in_mono_from(tcx, instance, set, target) - }); - } - } - /// If this basic block ends with a [`TerminatorKind::SwitchInt`] for which we can evaluate the /// dimscriminant in monomorphization, we return the discriminant bits and the /// [`SwitchTargets`], just so the caller doesn't also have to match on the terminator. diff --git a/compiler/rustc_middle/src/mir/traversal.rs b/compiler/rustc_middle/src/mir/traversal.rs index 0a938bcd31562..9321044c14b73 100644 --- a/compiler/rustc_middle/src/mir/traversal.rs +++ b/compiler/rustc_middle/src/mir/traversal.rs @@ -1,3 +1,5 @@ +use crate::query::Providers; + use super::*; /// Preorder traversal of a graph. @@ -279,3 +281,74 @@ pub fn reverse_postorder<'a, 'tcx>( { body.basic_blocks.reverse_postorder().iter().map(|&bb| (bb, &body.basic_blocks[bb])) } + +/// Finds which basic blocks are actually reachable for a monomorphized [`Instance`]. +/// +/// This is allowed to have false positives; just because this says a block +/// is reachable doesn't mean that's necessarily true. It's thus always +/// legal for this to return a filled set. +/// +/// Regardless, the [`BitSet::domain_size`] of the returned set will always +/// exactly match the number of blocks in the body so that `contains` +/// checks can be done without worrying about panicking. +/// +/// This is mostly useful because it lets us skip lowering the `false` side +/// of `if ::CONST`, as well as [`NullOp::UbChecks`]. +/// +/// [`NullOp::UbChecks`]: rustc_middle::mir::NullOp::UbChecks +fn reachable_blocks<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> BitSet { + let body = tcx.instance_mir(instance.def); + let mut visitor = MonoReachable { + body, + tcx, + instance, + worklist: Vec::new(), + visited: BitSet::new_empty(body.basic_blocks.len()), + }; + + visitor.visited.insert(START_BLOCK); + visitor.visit(START_BLOCK); + + while let Some(bb) = visitor.worklist.pop() { + if visitor.visited.insert(bb) { + visitor.visit(bb); + } + } + + visitor.visited +} + +struct MonoReachable<'a, 'tcx> { + body: &'a Body<'tcx>, + tcx: TyCtxt<'tcx>, + instance: Instance<'tcx>, + worklist: Vec, + visited: BitSet, +} + +impl<'a, 'tcx> MonoReachable<'a, 'tcx> { + fn visit(&mut self, bb: BasicBlock) { + let block = &self.body.basic_blocks[bb]; + + if let Some((bits, targets)) = + Body::try_const_mono_switchint(self.tcx, self.instance, block) + { + let target = targets.target_for_value(bits); + self.push(target); + } else { + for target in block.terminator().successors() { + self.push(target); + } + } + } + + fn push(&mut self, bb: BasicBlock) { + if !self.visited.contains(bb) { + self.worklist.push(bb); + } + } +} + +pub fn provide(providers: &mut Providers) { + providers.reachable_blocks = reachable_blocks; +} diff --git a/compiler/rustc_middle/src/query/keys.rs b/compiler/rustc_middle/src/query/keys.rs index 9cbc4d10146d1..96d2aef77381f 100644 --- a/compiler/rustc_middle/src/query/keys.rs +++ b/compiler/rustc_middle/src/query/keys.rs @@ -87,6 +87,15 @@ impl<'tcx> Key for ty::Instance<'tcx> { } } +impl<'tcx> AsLocalKey for ty::Instance<'tcx> { + type LocalKey = Self; + + #[inline(always)] + fn as_local_key(&self) -> Option { + self.def_id().is_local().then(|| *self) + } +} + impl<'tcx> Key for mir::interpret::GlobalId<'tcx> { type Cache = DefaultCache; @@ -534,6 +543,14 @@ impl<'tcx> Key for (ty::Instance<'tcx>, &'tcx ty::List>) { } } +impl<'tcx> Key for (ty::Instance<'tcx>, &'tcx mir::Body<'tcx>) { + type Cache = DefaultCache; + + fn default_span(&self, tcx: TyCtxt<'_>) -> Span { + self.0.default_span(tcx) + } +} + impl<'tcx> Key for (Ty<'tcx>, ty::ValTree<'tcx>) { type Cache = DefaultCache; diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs index 5e4454db3e28f..6c28d2442d4a1 100644 --- a/compiler/rustc_middle/src/query/mod.rs +++ b/compiler/rustc_middle/src/query/mod.rs @@ -71,6 +71,7 @@ use rustc_hir::def_id::{ }; use rustc_hir::lang_items::{LangItem, LanguageItems}; use rustc_hir::{Crate, ItemLocalId, ItemLocalMap, TraitCandidate}; +use rustc_index::bit_set::BitSet; use rustc_index::IndexVec; use rustc_query_system::ich::StableHashingContext; use rustc_query_system::query::{try_get_cached, QueryCache, QueryMode, QueryState}; @@ -270,7 +271,7 @@ rustc_queries! { feedable } - query unsizing_params_for_adt(key: DefId) -> &'tcx rustc_index::bit_set::BitSet + query unsizing_params_for_adt(key: DefId) -> &'tcx BitSet { arena_cache desc { |tcx| @@ -460,7 +461,7 @@ rustc_queries! { } /// Set of param indexes for type params that are in the type's representation - query params_in_repr(key: DefId) -> &'tcx rustc_index::bit_set::BitSet { + query params_in_repr(key: DefId) -> &'tcx BitSet { desc { "finding type parameters in the representation" } arena_cache no_hash @@ -2251,6 +2252,11 @@ rustc_queries! { query find_field((def_id, ident): (DefId, rustc_span::symbol::Ident)) -> Option { desc { |tcx| "find the index of maybe nested field `{ident}` in `{}`", tcx.def_path_str(def_id) } } + + query reachable_blocks(instance: ty::Instance<'tcx>) -> &'tcx BitSet { + arena_cache + desc { |tcx| "determining reachable blocks in `{}`", tcx.def_path_str(instance.def_id()) } + } } rustc_query_append! { define_callbacks! } diff --git a/compiler/rustc_monomorphize/Cargo.toml b/compiler/rustc_monomorphize/Cargo.toml index c7f1b9fa78454..3b99fd0728996 100644 --- a/compiler/rustc_monomorphize/Cargo.toml +++ b/compiler/rustc_monomorphize/Cargo.toml @@ -9,6 +9,7 @@ rustc_data_structures = { path = "../rustc_data_structures" } rustc_errors = { path = "../rustc_errors" } rustc_fluent_macro = { path = "../rustc_fluent_macro" } rustc_hir = { path = "../rustc_hir" } +rustc_index = { path = "../rustc_index" } rustc_macros = { path = "../rustc_macros" } rustc_middle = { path = "../rustc_middle" } rustc_session = { path = "../rustc_session" } diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs index 0c35f9838ed3f..396b5670e3f8d 100644 --- a/compiler/rustc_monomorphize/src/collector.rs +++ b/compiler/rustc_monomorphize/src/collector.rs @@ -211,6 +211,7 @@ use rustc_hir as hir; use rustc_hir::def::DefKind; use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId}; use rustc_hir::lang_items::LangItem; +use rustc_index::bit_set::BitSet; use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags; use rustc_middle::mir::interpret::{AllocId, ErrorHandled, GlobalAlloc, Scalar}; use rustc_middle::mir::mono::{InstantiationMode, MonoItem}; @@ -671,6 +672,7 @@ struct MirUsedCollector<'a, 'tcx> { visiting_call_terminator: bool, /// Set of functions for which it is OK to move large data into. skip_move_check_fns: Option>, + reachable_blocks: Option<&'tcx BitSet>, } impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> { @@ -831,6 +833,16 @@ impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> { } impl<'a, 'tcx> MirVisitor<'tcx> for MirUsedCollector<'a, 'tcx> { + fn visit_basic_block_data(&mut self, block: mir::BasicBlock, data: &mir::BasicBlockData<'tcx>) { + if self + .reachable_blocks + .expect("we should only walk blocks with CollectionMode::UsedItems") + .contains(block) + { + self.super_basic_block_data(block, data) + } + } + fn visit_rvalue(&mut self, rvalue: &mir::Rvalue<'tcx>, location: Location) { debug!("visiting rvalue {:?}", *rvalue); @@ -1411,6 +1423,11 @@ fn collect_items_of_instance<'tcx>( move_size_spans: vec![], visiting_call_terminator: false, skip_move_check_fns: None, + reachable_blocks: if mode == CollectionMode::UsedItems { + Some(tcx.reachable_blocks(instance)) + } else { + None + }, }; if mode == CollectionMode::UsedItems {