Skip to content

Commit

Permalink
Only collect mono items from reachable blocks
Browse files Browse the repository at this point in the history
  • Loading branch information
saethlin committed Mar 31, 2024
1 parent 395f780 commit 5fa7d14
Show file tree
Hide file tree
Showing 9 changed files with 122 additions and 59 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
Original file line number Diff line number Diff line change
Expand Up @@ -4374,6 +4374,7 @@ dependencies = [
"rustc_errors",
"rustc_fluent_macro",
"rustc_hir",
"rustc_index",
"rustc_macros",
"rustc_middle",
"rustc_session",
Expand Down
9 changes: 4 additions & 5 deletions compiler/rustc_codegen_ssa/src/mir/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
Expand Down
1 change: 1 addition & 0 deletions compiler/rustc_interface/src/passes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -631,6 +631,7 @@ pub static DEFAULT_QUERY_PROVIDERS: LazyLock<Providers> = LazyLock::new(|| {
rustc_lint::provide(providers);
rustc_symbol_mangling::provide(providers);
rustc_codegen_ssa::provide(providers);
rustc_middle::mir::traversal::provide(providers);
*providers
});

Expand Down
52 changes: 0 additions & 52 deletions compiler/rustc_middle/src/mir/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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 <T as Trait>::CONST`, as well as `intrinsics::debug_assertions`.
pub fn reachable_blocks_in_mono(
&self,
tcx: TyCtxt<'tcx>,
instance: Instance<'tcx>,
) -> BitSet<BasicBlock> {
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<BasicBlock>,
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.
Expand Down
73 changes: 73 additions & 0 deletions compiler/rustc_middle/src/mir/traversal.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use crate::query::Providers;

use super::*;

/// Preorder traversal of a graph.
Expand Down Expand Up @@ -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 <T as Trait>::CONST`, as well as [`NullOp::UbChecks`].
///
/// [`NullOp::UbChecks`]: rustc_middle::mir::NullOp::UbChecks
fn reachable_blocks<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> BitSet<BasicBlock> {
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<BasicBlock>,
visited: BitSet<BasicBlock>,
}

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;
}
17 changes: 17 additions & 0 deletions compiler/rustc_middle/src/query/keys.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::LocalKey> {
self.def_id().is_local().then(|| *self)
}
}

impl<'tcx> Key for mir::interpret::GlobalId<'tcx> {
type Cache<V> = DefaultCache<Self, V>;

Expand Down Expand Up @@ -534,6 +543,14 @@ impl<'tcx> Key for (ty::Instance<'tcx>, &'tcx ty::List<Ty<'tcx>>) {
}
}

impl<'tcx> Key for (ty::Instance<'tcx>, &'tcx mir::Body<'tcx>) {
type Cache<V> = DefaultCache<Self, V>;

fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
self.0.default_span(tcx)
}
}

impl<'tcx> Key for (Ty<'tcx>, ty::ValTree<'tcx>) {
type Cache<V> = DefaultCache<Self, V>;

Expand Down
10 changes: 8 additions & 2 deletions compiler/rustc_middle/src/query/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -270,7 +271,7 @@ rustc_queries! {
feedable
}

query unsizing_params_for_adt(key: DefId) -> &'tcx rustc_index::bit_set::BitSet<u32>
query unsizing_params_for_adt(key: DefId) -> &'tcx BitSet<u32>
{
arena_cache
desc { |tcx|
Expand Down Expand Up @@ -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<u32> {
query params_in_repr(key: DefId) -> &'tcx BitSet<u32> {
desc { "finding type parameters in the representation" }
arena_cache
no_hash
Expand Down Expand Up @@ -2251,6 +2252,11 @@ rustc_queries! {
query find_field((def_id, ident): (DefId, rustc_span::symbol::Ident)) -> Option<rustc_target::abi::FieldIdx> {
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<mir::BasicBlock> {
arena_cache
desc { |tcx| "determining reachable blocks in `{}`", tcx.def_path_str(instance.def_id()) }
}
}

rustc_query_append! { define_callbacks! }
Expand Down
1 change: 1 addition & 0 deletions compiler/rustc_monomorphize/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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" }
Expand Down
17 changes: 17 additions & 0 deletions compiler/rustc_monomorphize/src/collector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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<Vec<DefId>>,
reachable_blocks: Option<&'tcx BitSet<mir::BasicBlock>>,
}

impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 5fa7d14

Please sign in to comment.