Skip to content

Commit 9fea3d8

Browse files
scottmcmscottmcm
committed
[MIR-OPT] Add a pass that replicates simple branches into predecessors
This means fewer blocks, so is hopefully a net win.
1 parent c2cabee commit 9fea3d8

10 files changed

+503
-119
lines changed

‎compiler/rustc_mir_transform/src/lib.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ mod remove_storage_markers;
9090
mod remove_uninit_drops;
9191
mod remove_unneeded_drops;
9292
mod remove_zsts;
93+
mod replicate_branches;
9394
mod required_consts;
9495
mod reveal_all;
9596
mod separate_const_switch;
@@ -554,7 +555,8 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
554555
&remove_unneeded_drops::RemoveUnneededDrops,
555556
&sroa::ScalarReplacementOfAggregates,
556557
&match_branches::MatchBranchSimplification,
557-
// inst combine is after MatchBranchSimplification to clean up Ne(_1, false)
558+
// inst simplify is after MatchBranchSimplification to clean up Ne(_1, false)
559+
&replicate_branches::ReplicateBranches,
558560
&multiple_return_terminators::MultipleReturnTerminators,
559561
&instsimplify::InstSimplify,
560562
&separate_const_switch::SeparateConstSwitch,

‎compiler/rustc_mir_transform/src/replicate_branches.rs

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
//! Looks for basic blocks that are just a simple 2-way branch,
2+
//! and replicates them into their predecessors.
3+
//!
4+
//! That sounds wasteful, but it's very common that the predecessor just set
5+
//! whatever we're about to branch on, and this makes that easier to collapse
6+
//! later. And if not, well, an extra branch is no big deal.
7+
8+
use crate::MirPass;
9+
10+
use rustc_middle::mir::*;
11+
use rustc_middle::ty::TyCtxt;
12+
13+
use super::simplify::simplify_cfg;
14+
15+
pub struct ReplicateBranches;
16+
17+
impl<'tcx> MirPass<'tcx> for ReplicateBranches {
18+
fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
19+
sess.mir_opt_level() >= 2
20+
}
21+
22+
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
23+
debug!("Running ReplicateBranches on `{:?}`", body.source);
24+
25+
let mut source_and_target_blocks = vec![];
26+
27+
for (bb, data) in body.basic_blocks.iter_enumerated() {
28+
if data.is_cleanup {
29+
continue;
30+
}
31+
32+
let TerminatorKind::SwitchInt { discr, targets } = &data.terminator().kind
33+
else { continue };
34+
35+
let Some(place) = discr.place() else { continue };
36+
37+
// Only replicate simple branches. That means either:
38+
// - A specific target plus a potentially-reachable otherwise, or
39+
// - Two specific targets and the otherwise is unreachable.
40+
if !(targets.iter().len() < 2
41+
|| (targets.iter().len() <= 2
42+
&& body.basic_blocks[targets.otherwise()].is_empty_unreachable()))
43+
{
44+
continue;
45+
}
46+
47+
// Only replicate blocks that branch on a discriminant.
48+
let mut found_discriminant = false;
49+
for stmt in &data.statements {
50+
match &stmt.kind {
51+
StatementKind::StorageDead { .. } | StatementKind::StorageLive { .. } => {}
52+
StatementKind::Assign(place_and_rvalue)
53+
if place_and_rvalue.0 == place
54+
&& matches!(place_and_rvalue.1, Rvalue::Discriminant(..)) =>
55+
{
56+
found_discriminant = true;
57+
}
58+
_ => continue,
59+
}
60+
}
61+
62+
// This currently doesn't duplicate ordinary boolean checks as that regresses
63+
// various tests that seem to depend on the existing structure.
64+
if !found_discriminant {
65+
continue;
66+
}
67+
68+
// Only replicate to a small number of `goto` predecessors.
69+
let preds = &body.basic_blocks.predecessors()[bb];
70+
if preds.len() > 2
71+
|| !preds.iter().copied().all(|p| {
72+
let pdata = &body.basic_blocks[p];
73+
matches!(pdata.terminator().kind, TerminatorKind::Goto { .. })
74+
&& !pdata.is_cleanup
75+
})
76+
{
77+
continue;
78+
}
79+
80+
for p in preds {
81+
source_and_target_blocks.push((*p, bb));
82+
}
83+
}
84+
85+
if source_and_target_blocks.is_empty() {
86+
return;
87+
}
88+
89+
let basic_blocks = body.basic_blocks.as_mut();
90+
for (source, target) in source_and_target_blocks {
91+
debug_assert_eq!(
92+
basic_blocks[source].terminator().kind,
93+
TerminatorKind::Goto { target },
94+
);
95+
96+
let (source, target) = basic_blocks.pick2_mut(source, target);
97+
source.statements.extend(target.statements.iter().cloned());
98+
source.terminator = target.terminator.clone();
99+
}
100+
101+
// The target blocks should be unused now, so cleanup after ourselves.
102+
simplify_cfg(tcx, body);
103+
}
104+
}

‎tests/mir-opt/inline/unchecked_shifts.unchecked_shl_unsigned_smaller.Inline.diff

Lines changed: 90 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,51 @@
1010
+ scope 1 (inlined core::num::<impl u16>::unchecked_shl) { // at $DIR/unchecked_shifts.rs:11:7: 11:23
1111
+ debug self => _3; // in scope 1 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
1212
+ debug rhs => _4; // in scope 1 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
13-
+ let mut _5: u16; // in scope 1 at $SRC_DIR/core/src/num/mod.rs:LL:COL
14-
+ let mut _6: (u32,); // in scope 1 at $SRC_DIR/core/src/num/mod.rs:LL:COL
13+
+ let mut _5: (u32,); // in scope 1 at $SRC_DIR/core/src/num/mod.rs:LL:COL
14+
+ let mut _6: u32; // in scope 1 at $SRC_DIR/core/src/num/mod.rs:LL:COL
1515
+ scope 2 {
16+
+ scope 3 (inlined core::num::<impl u16>::unchecked_shl::conv) { // at $SRC_DIR/core/src/num/mod.rs:LL:COL
17+
+ debug x => _6; // in scope 3 at $SRC_DIR/core/src/num/mod.rs:LL:COL
18+
+ let mut _7: std::option::Option<u16>; // in scope 3 at $SRC_DIR/core/src/num/mod.rs:LL:COL
19+
+ let mut _8: std::result::Result<u16, std::num::TryFromIntError>; // in scope 3 at $SRC_DIR/core/src/num/mod.rs:LL:COL
20+
+ scope 4 {
21+
+ scope 5 (inlined <u32 as TryInto<u16>>::try_into) { // at $SRC_DIR/core/src/num/mod.rs:LL:COL
22+
+ debug self => _6; // in scope 5 at $SRC_DIR/core/src/convert/mod.rs:LL:COL
23+
+ scope 6 (inlined convert::num::<impl TryFrom<u32> for u16>::try_from) { // at $SRC_DIR/core/src/convert/mod.rs:LL:COL
24+
+ debug u => _6; // in scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
25+
+ let mut _9: bool; // in scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
26+
+ let mut _10: u32; // in scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
27+
+ let mut _11: u16; // in scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
28+
+ }
29+
+ }
30+
+ scope 7 (inlined Result::<u16, TryFromIntError>::ok) { // at $SRC_DIR/core/src/num/mod.rs:LL:COL
31+
+ debug self => _8; // in scope 7 at $SRC_DIR/core/src/result.rs:LL:COL
32+
+ let _12: u16; // in scope 7 at $SRC_DIR/core/src/result.rs:LL:COL
33+
+ scope 8 {
34+
+ debug x => _12; // in scope 8 at $SRC_DIR/core/src/result.rs:LL:COL
35+
+ }
36+
+ }
37+
+ scope 9 (inlined #[track_caller] Option::<u16>::unwrap_unchecked) { // at $SRC_DIR/core/src/num/mod.rs:LL:COL
38+
+ debug self => _7; // in scope 9 at $SRC_DIR/core/src/option.rs:LL:COL
39+
+ let mut _13: &std::option::Option<u16>; // in scope 9 at $SRC_DIR/core/src/option.rs:LL:COL
40+
+ let _14: u16; // in scope 9 at $SRC_DIR/core/src/option.rs:LL:COL
41+
+ scope 10 {
42+
+ debug val => _14; // in scope 10 at $SRC_DIR/core/src/option.rs:LL:COL
43+
+ }
44+
+ scope 11 {
45+
+ scope 13 (inlined unreachable_unchecked) { // at $SRC_DIR/core/src/option.rs:LL:COL
46+
+ scope 14 {
47+
+ scope 15 (inlined unreachable_unchecked::runtime) { // at $SRC_DIR/core/src/intrinsics.rs:LL:COL
48+
+ }
49+
+ }
50+
+ }
51+
+ }
52+
+ scope 12 (inlined Option::<u16>::is_some) { // at $SRC_DIR/core/src/option.rs:LL:COL
53+
+ debug self => _13; // in scope 12 at $SRC_DIR/core/src/option.rs:LL:COL
54+
+ }
55+
+ }
56+
+ }
57+
+ }
1658
+ }
1759
+ }
1860

@@ -22,30 +64,61 @@
2264
StorageLive(_4); // scope 0 at $DIR/unchecked_shifts.rs:+1:21: +1:22
2365
_4 = _2; // scope 0 at $DIR/unchecked_shifts.rs:+1:21: +1:22
2466
- _0 = core::num::<impl u16>::unchecked_shl(move _3, move _4) -> bb1; // scope 0 at $DIR/unchecked_shifts.rs:+1:5: +1:23
25-
+ StorageLive(_5); // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
26-
+ StorageLive(_6); // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
27-
+ _6 = (_4,); // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
28-
+ _5 = core::num::<impl u16>::unchecked_shl::conv(move (_6.0: u32)) -> bb1; // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
29-
// mir::Constant
67+
- // mir::Constant
3068
- // + span: $DIR/unchecked_shifts.rs:11:7: 11:20
3169
- // + literal: Const { ty: unsafe fn(u16, u32) -> u16 {core::num::<impl u16>::unchecked_shl}, val: Value(<ZST>) }
32-
+ // + span: $SRC_DIR/core/src/num/mod.rs:LL:COL
33-
+ // + literal: Const { ty: fn(u32) -> u16 {core::num::<impl u16>::unchecked_shl::conv}, val: Value(<ZST>) }
70+
+ StorageLive(_14); // scope 0 at $DIR/unchecked_shifts.rs:+1:7: +1:23
71+
+ StorageLive(_5); // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
72+
+ _5 = (_4,); // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
73+
+ StorageLive(_6); // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
74+
+ _6 = move (_5.0: u32); // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
75+
+ StorageLive(_7); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
76+
+ StorageLive(_8); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
77+
+ StorageLive(_9); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
78+
+ StorageLive(_10); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
79+
+ _10 = const 65535_u32; // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
80+
+ _9 = Gt(_6, move _10); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
81+
+ StorageDead(_10); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
82+
+ switchInt(move _9) -> [0: bb3, otherwise: bb2]; // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
3483
}
3584

3685
bb1: {
86+
+ StorageDead(_14); // scope 0 at $DIR/unchecked_shifts.rs:+1:7: +1:23
87+
StorageDead(_4); // scope 0 at $DIR/unchecked_shifts.rs:+1:22: +1:23
88+
StorageDead(_3); // scope 0 at $DIR/unchecked_shifts.rs:+1:22: +1:23
89+
return; // scope 0 at $DIR/unchecked_shifts.rs:+2:2: +2:2
90+
+ }
91+
+
92+
+ bb2: {
93+
+ StorageDead(_9); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
94+
+ StorageLive(_12); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
95+
+ StorageDead(_12); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
96+
+ StorageDead(_8); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
97+
+ StorageLive(_13); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
98+
+ unreachable; // scope 7 at $SRC_DIR/core/src/result.rs:LL:COL
99+
+ }
100+
+
101+
+ bb3: {
102+
+ StorageLive(_11); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
103+
+ _11 = _6 as u16 (IntToInt); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
104+
+ _8 = Result::<u16, TryFromIntError>::Ok(move _11); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
105+
+ StorageDead(_11); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
106+
+ StorageDead(_9); // scope 6 at $SRC_DIR/core/src/convert/num.rs:LL:COL
107+
+ StorageLive(_12); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
108+
+ _12 = move ((_8 as Ok).0: u16); // scope 7 at $SRC_DIR/core/src/result.rs:LL:COL
109+
+ _7 = Option::<u16>::Some(move _12); // scope 8 at $SRC_DIR/core/src/result.rs:LL:COL
110+
+ StorageDead(_12); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
111+
+ StorageDead(_8); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
112+
+ StorageLive(_13); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
113+
+ _14 = move ((_7 as Some).0: u16); // scope 9 at $SRC_DIR/core/src/option.rs:LL:COL
114+
+ StorageDead(_13); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
115+
+ StorageDead(_7); // scope 4 at $SRC_DIR/core/src/num/mod.rs:LL:COL
37116
+ StorageDead(_6); // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
38-
+ _0 = unchecked_shl::<u16>(_3, move _5) -> [return: bb2, unwind unreachable]; // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
117+
+ StorageDead(_5); // scope 2 at $SRC_DIR/core/src/num/mod.rs:LL:COL
118+
+ _0 = unchecked_shl::<u16>(_3, move _14) -> [return: bb1, unwind unreachable]; // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
39119
+ // mir::Constant
40120
+ // + span: $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
41121
+ // + literal: Const { ty: unsafe extern "rust-intrinsic" fn(u16, u16) -> u16 {unchecked_shl::<u16>}, val: Value(<ZST>) }
42-
+ }
43-
+
44-
+ bb2: {
45-
+ StorageDead(_5); // scope 2 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
46-
StorageDead(_4); // scope 0 at $DIR/unchecked_shifts.rs:+1:22: +1:23
47-
StorageDead(_3); // scope 0 at $DIR/unchecked_shifts.rs:+1:22: +1:23
48-
return; // scope 0 at $DIR/unchecked_shifts.rs:+2:2: +2:2
49122
}
50123
}
51124

0 commit comments

Comments
 (0)