Skip to content

Commit b39e4ca

Browse files
borsbors
committed
Auto merge of #121421 - saethlin:smarter-mono, r=<try>
Avoid lowering code under dead SwitchInt targets r? `@ghost` Reviving #91222 per #120848
2 parents 3406ada + 6b3d152 commit b39e4ca

13 files changed

+512
-334
lines changed

‎compiler/rustc_codegen_ssa/src/mir/block.rs

+10
Original file line numberDiff line numberDiff line change
@@ -1179,6 +1179,16 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
11791179
}
11801180
}
11811181

1182+
pub fn codegen_block_as_unreachable(&mut self, bb: mir::BasicBlock) {
1183+
let llbb = match self.try_llbb(bb) {
1184+
Some(llbb) => llbb,
1185+
None => return,
1186+
};
1187+
let bx = &mut Bx::build(self.cx, llbb);
1188+
debug!("codegen_block_as_unreachable({:?})", bb);
1189+
bx.unreachable();
1190+
}
1191+
11821192
fn codegen_terminator(
11831193
&mut self,
11841194
bx: &mut Bx,

‎compiler/rustc_codegen_ssa/src/mir/mod.rs

+10-1
Original file line numberDiff line numberDiff line change
@@ -256,13 +256,22 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
256256
// Apply debuginfo to the newly allocated locals.
257257
fx.debug_introduce_locals(&mut start_bx);
258258

259+
let reachable_blocks = mir.reachable_blocks_in_mono(cx.tcx(), instance);
260+
259261
// The builders will be created separately for each basic block at `codegen_block`.
260262
// So drop the builder of `start_llbb` to avoid having two at the same time.
261263
drop(start_bx);
262264

263265
// Codegen the body of each block using reverse postorder
264266
for (bb, _) in traversal::reverse_postorder(mir) {
265-
fx.codegen_block(bb);
267+
if reachable_blocks.contains(bb) {
268+
fx.codegen_block(bb);
269+
} else {
270+
// This may have references to things we didn't monomorphize, so we
271+
// don't actually codegen the body. We still create the block so
272+
// terminators in other blocks can reference it without worry.
273+
fx.codegen_block_as_unreachable(bb);
274+
}
266275
}
267276
}
268277

‎compiler/rustc_middle/src/mir/mod.rs

+69-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use crate::ty::print::{pretty_print_const, with_no_trimmed_paths};
1010
use crate::ty::print::{FmtPrinter, Printer};
1111
use crate::ty::visit::TypeVisitableExt;
1212
use crate::ty::{self, List, Ty, TyCtxt};
13-
use crate::ty::{AdtDef, InstanceDef, UserTypeAnnotationIndex};
13+
use crate::ty::{AdtDef, Instance, InstanceDef, UserTypeAnnotationIndex};
1414
use crate::ty::{GenericArg, GenericArgsRef};
1515

1616
use rustc_data_structures::captures::Captures;
@@ -29,6 +29,7 @@ pub use rustc_ast::Mutability;
2929
use rustc_data_structures::fx::FxHashMap;
3030
use rustc_data_structures::fx::FxHashSet;
3131
use rustc_data_structures::graph::dominators::Dominators;
32+
use rustc_index::bit_set::BitSet;
3233
use rustc_index::{Idx, IndexSlice, IndexVec};
3334
use rustc_serialize::{Decodable, Encodable};
3435
use rustc_span::symbol::Symbol;
@@ -642,6 +643,73 @@ impl<'tcx> Body<'tcx> {
642643
self.injection_phase.is_some()
643644
}
644645

646+
/// Finds which basic blocks are actually reachable for a specific
647+
/// monomorphization of this body.
648+
///
649+
/// This is allowed to have false positives; just because this says a block
650+
/// is reachable doesn't mean that's necessarily true. It's thus always
651+
/// legal for this to return a filled set.
652+
///
653+
/// Regardless, the [`BitSet::domain_size`] of the returned set will always
654+
/// exactly match the number of blocks in the body so that `contains`
655+
/// checks can be done without worrying about panicking.
656+
///
657+
/// The main case this supports is filtering out `if <T as Trait>::CONST`
658+
/// bodies that can't be removed in generic MIR, but *can* be removed once
659+
/// the specific `T` is known.
660+
///
661+
/// This is used in the monomorphization collector as well as in codegen.
662+
pub fn reachable_blocks_in_mono(
663+
&self,
664+
tcx: TyCtxt<'tcx>,
665+
instance: Instance<'tcx>,
666+
) -> BitSet<BasicBlock> {
667+
if instance.args.non_erasable_generics(tcx, instance.def_id()).next().is_none() {
668+
// If it's non-generic, then mir-opt const prop has already run, meaning it's
669+
// probably not worth doing any further filtering. So call everything reachable.
670+
return BitSet::new_filled(self.basic_blocks.len());
671+
}
672+
673+
let mut set = BitSet::new_empty(self.basic_blocks.len());
674+
self.reachable_blocks_in_mono_from(tcx, instance, &mut set, START_BLOCK);
675+
set
676+
}
677+
678+
fn reachable_blocks_in_mono_from(
679+
&self,
680+
tcx: TyCtxt<'tcx>,
681+
instance: Instance<'tcx>,
682+
set: &mut BitSet<BasicBlock>,
683+
bb: BasicBlock,
684+
) {
685+
if !set.insert(bb) {
686+
return;
687+
}
688+
689+
let data = &self.basic_blocks[bb];
690+
691+
if let TerminatorKind::SwitchInt { discr: Operand::Constant(constant), targets } =
692+
&data.terminator().kind
693+
{
694+
let env = ty::ParamEnv::reveal_all();
695+
let mono_literal = instance.instantiate_mir_and_normalize_erasing_regions(
696+
tcx,
697+
env,
698+
crate::ty::EarlyBinder::bind(constant.const_),
699+
);
700+
if let Some(bits) = mono_literal.try_eval_bits(tcx, env) {
701+
let target = targets.target_for_value(bits);
702+
return self.reachable_blocks_in_mono_from(tcx, instance, set, target);
703+
} else {
704+
bug!("Couldn't evaluate constant {:?} in mono {:?}", constant, instance);
705+
}
706+
}
707+
708+
for target in data.terminator().successors() {
709+
self.reachable_blocks_in_mono_from(tcx, instance, set, target);
710+
}
711+
}
712+
645713
/// For a `Location` in this scope, determine what the "caller location" at that point is. This
646714
/// is interesting because of inlining: the `#[track_caller]` attribute of inlined functions
647715
/// must be honored. Falls back to the `tracked_caller` value for `#[track_caller]` functions,

‎compiler/rustc_middle/src/mir/traversal.rs

-2
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
use rustc_index::bit_set::BitSet;
2-
31
use super::*;
42

53
/// Preorder traversal of a graph.

‎compiler/rustc_mir_transform/src/lib.rs

+2
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ mod check_alignment;
108108
pub mod simplify;
109109
mod simplify_branches;
110110
mod simplify_comparison_integral;
111+
mod simplify_if_const;
111112
mod sroa;
112113
mod uninhabited_enum_branching;
113114
mod unreachable_prop;
@@ -616,6 +617,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
616617
&large_enums::EnumSizeOpt { discrepancy: 128 },
617618
// Some cleanup necessary at least for LLVM and potentially other codegen backends.
618619
&add_call_guards::CriticalCallEdges,
620+
&simplify_if_const::SimplifyIfConst,
619621
// Cleanup for human readability, off by default.
620622
&prettify::ReorderBasicBlocks,
621623
&prettify::ReorderLocals,

‎compiler/rustc_mir_transform/src/simplify_if_const.rs

+76
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
//! A pass that simplifies branches when their condition is known.
2+
3+
use crate::MirPass;
4+
use rustc_middle::mir::*;
5+
use rustc_middle::ty::TyCtxt;
6+
7+
/// The lowering for `if CONST` produces
8+
/// ```
9+
/// _1 = Const(...);
10+
/// switchInt (move _1)
11+
/// ```
12+
/// so this pass replaces that with
13+
/// ```
14+
/// switchInt (Const(...))
15+
/// ```
16+
/// so that further MIR consumers can special-case it more easily.
17+
///
18+
/// Unlike ConstProp, this supports generic constants too, not just concrete ones.
19+
pub struct SimplifyIfConst;
20+
21+
impl<'tcx> MirPass<'tcx> for SimplifyIfConst {
22+
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
23+
for block in body.basic_blocks_mut() {
24+
simplify_assign_move_switch(tcx, block);
25+
}
26+
}
27+
}
28+
29+
fn simplify_assign_move_switch(tcx: TyCtxt<'_>, block: &mut BasicBlockData<'_>) {
30+
let Some(Terminator { kind: TerminatorKind::SwitchInt { discr: switch_desc, .. }, .. }) =
31+
&mut block.terminator
32+
else {
33+
return;
34+
};
35+
36+
let &mut Operand::Move(switch_place) = &mut *switch_desc else { return };
37+
38+
let Some(switch_local) = switch_place.as_local() else { return };
39+
40+
let Some(last_statement) = block.statements.last_mut() else { return };
41+
42+
let StatementKind::Assign(boxed_place_rvalue) = &last_statement.kind else { return };
43+
44+
let Some(assigned_local) = boxed_place_rvalue.0.as_local() else { return };
45+
46+
if switch_local != assigned_local {
47+
return;
48+
}
49+
50+
if !matches!(boxed_place_rvalue.1, Rvalue::Use(Operand::Constant(_))) {
51+
return;
52+
}
53+
54+
let should_optimize = tcx.consider_optimizing(|| {
55+
format!(
56+
"SimplifyBranches - Assignment: {:?} SourceInfo: {:?}",
57+
boxed_place_rvalue, last_statement.source_info
58+
)
59+
});
60+
61+
if should_optimize {
62+
let Some(last_statement) = block.statements.pop() else {
63+
bug!("Somehow the statement disappeared?");
64+
};
65+
66+
let StatementKind::Assign(boxed_place_rvalue) = last_statement.kind else {
67+
bug!("Somehow it's not an assignment any more?");
68+
};
69+
70+
let Rvalue::Use(assigned_constant @ Operand::Constant(_)) = boxed_place_rvalue.1 else {
71+
bug!("Somehow it's not a use of a constant any more?");
72+
};
73+
74+
*switch_desc = assigned_constant;
75+
}
76+
}

‎tests/codegen/skip-mono-inside-if-false.rs

+39
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
// compile-flags: -O -C no-prepopulate-passes
2+
3+
#![crate_type = "lib"]
4+
5+
#[no_mangle]
6+
pub fn demo_for_i32() {
7+
generic_impl::<i32>();
8+
}
9+
10+
// Two important things here:
11+
// - We replace the "then" block with `unreachable` to avoid linking problems
12+
// - We neither declare nor define the `big_impl` that said block "calls".
13+
14+
// CHECK-LABEL: ; skip_mono_inside_if_false::generic_impl
15+
// CHECK: start:
16+
// CHECK-NEXT: br i1 false, label %[[THEN_BRANCH:bb[0-9]+]], label %[[ELSE_BRANCH:bb[0-9]+]]
17+
// CHECK: [[ELSE_BRANCH]]:
18+
// CHECK-NEXT: call skip_mono_inside_if_false::small_impl
19+
// CHECK: [[THEN_BRANCH]]:
20+
// CHECK-NEXT: unreachable
21+
22+
fn generic_impl<T>() {
23+
trait MagicTrait {
24+
const IS_BIG: bool;
25+
}
26+
impl<T> MagicTrait for T {
27+
const IS_BIG: bool = std::mem::size_of::<T>() > 10;
28+
}
29+
if T::IS_BIG {
30+
big_impl::<T>();
31+
} else {
32+
small_impl::<T>();
33+
}
34+
}
35+
36+
#[inline(never)]
37+
fn small_impl<T>() {}
38+
#[inline(never)]
39+
fn big_impl<T>() {}

0 commit comments

Comments
 (0)