Skip to content

Commit 2dd720c

Browse files
scottmcmscottmcm
committed
Use multiple returns in MIR if it saves a block; still have only one in LLVM
1 parent 30f168e commit 2dd720c

File tree

37 files changed

+310
-201
lines changed

37 files changed

+310
-201
lines changed

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

+10-2
Original file line numberDiff line numberDiff line change
@@ -462,7 +462,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
462462
}
463463
}
464464

465-
fn codegen_return_terminator(&mut self, bx: &mut Bx) {
465+
pub(super) fn codegen_return_terminator(&mut self, bx: &mut Bx) {
466466
// Call `va_end` if this is the definition of a C-variadic function.
467467
if self.fn_abi.c_variadic {
468468
// The `VaList` "spoofed" argument is just after all the real arguments.
@@ -1343,7 +1343,15 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
13431343
}
13441344

13451345
mir::TerminatorKind::Return => {
1346-
self.codegen_return_terminator(bx);
1346+
match self.return_block {
1347+
CachedLlbb::Skip => self.codegen_return_terminator(bx),
1348+
CachedLlbb::Some(target) => bx.br(target),
1349+
CachedLlbb::None => {
1350+
let return_llbb = bx.append_sibling_block("return");
1351+
self.return_block = CachedLlbb::Some(return_llbb);
1352+
bx.br(return_llbb);
1353+
}
1354+
}
13471355
MergingSucc::False
13481356
}
13491357

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

+24
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,12 @@ pub struct FunctionCx<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> {
9292
/// Cached terminate upon unwinding block and its reason
9393
terminate_block: Option<(Bx::BasicBlock, UnwindTerminateReason)>,
9494

95+
/// Shared return block, because LLVM would prefer only one `ret`.
96+
///
97+
/// If this is `Skip`, there's only one return in the function (or none at all)
98+
/// so there's no shared return, just the one in the normal BB.
99+
return_block: CachedLlbb<Bx::BasicBlock>,
100+
95101
/// A bool flag for each basic block indicating whether it is a cold block.
96102
/// A cold block is a block that is unlikely to be executed at runtime.
97103
cold_blocks: IndexVec<mir::BasicBlock, bool>,
@@ -204,6 +210,18 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
204210
})
205211
.collect();
206212

213+
let return_block = if mir
214+
.basic_blocks
215+
.iter()
216+
.filter(|bbd| matches!(bbd.terminator().kind, mir::TerminatorKind::Return))
217+
.count()
218+
> 1
219+
{
220+
CachedLlbb::None
221+
} else {
222+
CachedLlbb::Skip
223+
};
224+
207225
let mut fx = FunctionCx {
208226
instance,
209227
mir,
@@ -214,6 +232,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
214232
cached_llbbs,
215233
unreachable_block: None,
216234
terminate_block: None,
235+
return_block,
217236
cleanup_kinds,
218237
landing_pads: IndexVec::from_elem(None, &mir.basic_blocks),
219238
funclets: IndexVec::from_fn_n(|_| None, mir.basic_blocks.len()),
@@ -308,6 +327,11 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
308327
for bb in unreached_blocks.iter() {
309328
fx.codegen_block_as_unreachable(bb);
310329
}
330+
331+
if let CachedLlbb::Some(llbb) = fx.return_block {
332+
let bx = &mut Bx::build(fx.cx, llbb);
333+
fx.codegen_return_terminator(bx);
334+
}
311335
}
312336

313337
/// Produces, for each argument, a `Value` pointing at the

‎compiler/rustc_mir_transform/src/lib.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -694,7 +694,6 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
694694
&sroa::ScalarReplacementOfAggregates,
695695
&match_branches::MatchBranchSimplification,
696696
// inst combine is after MatchBranchSimplification to clean up Ne(_1, false)
697-
&multiple_return_terminators::MultipleReturnTerminators,
698697
// After simplifycfg, it allows us to discover new opportunities for peephole
699698
// optimizations.
700699
&instsimplify::InstSimplify::AfterSimplifyCfg,
@@ -711,14 +710,15 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
711710
&dest_prop::DestinationPropagation,
712711
&o1(simplify_branches::SimplifyConstCondition::Final),
713712
&o1(remove_noop_landing_pads::RemoveNoopLandingPads),
713+
// Can make blocks unused, so before the last simplify-cfg
714+
&multiple_return_terminators::MultipleReturnTerminators,
714715
&o1(simplify::SimplifyCfg::Final),
715716
// After the last SimplifyCfg, because this wants one-block functions.
716717
&strip_debuginfo::StripDebugInfo,
717718
&copy_prop::CopyProp,
718719
&dead_store_elimination::DeadStoreElimination::Final,
719720
&nrvo::RenameReturnPlace,
720721
&simplify::SimplifyLocals::Final,
721-
&multiple_return_terminators::MultipleReturnTerminators,
722722
&large_enums::EnumSizeOpt { discrepancy: 128 },
723723
// Some cleanup necessary at least for LLVM and potentially other codegen backends.
724724
&add_call_guards::CriticalCallEdges,

‎compiler/rustc_mir_transform/src/multiple_return_terminators.rs

+29-18
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,51 @@
11
//! This pass removes jumps to basic blocks containing only a return, and replaces them with a
22
//! return instead.
33
4-
use rustc_index::bit_set::DenseBitSet;
54
use rustc_middle::mir::*;
65
use rustc_middle::ty::TyCtxt;
7-
8-
use crate::simplify;
6+
use smallvec::SmallVec;
97

108
pub(super) struct MultipleReturnTerminators;
119

1210
impl<'tcx> crate::MirPass<'tcx> for MultipleReturnTerminators {
1311
fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
14-
sess.mir_opt_level() >= 4
12+
sess.mir_opt_level() >= 2
1513
}
1614

1715
fn run_pass(&self, _: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
18-
// find basic blocks with no statement and a return terminator
19-
let mut bbs_simple_returns = DenseBitSet::new_empty(body.basic_blocks.len());
20-
let bbs = body.basic_blocks_mut();
21-
for idx in bbs.indices() {
22-
if bbs[idx].statements.is_empty()
23-
&& bbs[idx].terminator().kind == TerminatorKind::Return
16+
let mut to_handle = <Vec<(BasicBlock, SmallVec<_>)>>::new();
17+
for (bb, bbdata) in body.basic_blocks.iter_enumerated() {
18+
// Look for returns where, if we lift them into the parents, we can save a block.
19+
if let TerminatorKind::Return = bbdata.terminator().kind
20+
&& bbdata
21+
.statements
22+
.iter()
23+
.all(|stmt| matches!(stmt.kind, StatementKind::StorageDead(_)))
24+
&& let predecessors = &body.basic_blocks.predecessors()[bb]
25+
&& predecessors.len() >= 2
26+
&& predecessors.iter().all(|pred| {
27+
matches!(
28+
body.basic_blocks[*pred].terminator().kind,
29+
TerminatorKind::Goto { .. },
30+
)
31+
})
2432
{
25-
bbs_simple_returns.insert(idx);
33+
to_handle.push((bb, predecessors.clone()));
2634
}
2735
}
2836

29-
for bb in bbs {
30-
if let TerminatorKind::Goto { target } = bb.terminator().kind {
31-
if bbs_simple_returns.contains(target) {
32-
bb.terminator_mut().kind = TerminatorKind::Return;
33-
}
34-
}
37+
if to_handle.is_empty() {
38+
return;
3539
}
3640

37-
simplify::remove_dead_blocks(body)
41+
let bbs = body.basic_blocks_mut();
42+
for (succ, predecessors) in to_handle {
43+
for pred in predecessors {
44+
let (pred_block, succ_block) = bbs.pick2_mut(pred, succ);
45+
pred_block.statements.extend(succ_block.statements.iter().cloned());
46+
*pred_block.terminator_mut() = succ_block.terminator().clone();
47+
}
48+
}
3849
}
3950

4051
fn is_required(&self) -> bool {

‎tests/codegen/asm/goto.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ pub unsafe fn asm_goto() {
1919
pub unsafe fn asm_goto_with_outputs() -> u64 {
2020
let out: u64;
2121
// CHECK: [[RES:%[0-9]+]] = callbr i64 asm sideeffect alignstack inteldialect "
22-
// CHECK-NEXT: to label %[[FALLTHROUGHBB:[a-b0-9]+]] [label %[[JUMPBB:[a-b0-9]+]]]
22+
// CHECK-NEXT: to label %[[FALLTHROUGHBB:[a-b0-9]+]] [label %[[JUMPBB:return]]]
2323
asm!("{} /* {} */", out(reg) out, label { return 1; });
2424
// CHECK: [[JUMPBB]]:
2525
// CHECK-NEXT: [[RET:%.+]] = phi i64 [ [[RES]], %[[FALLTHROUGHBB]] ], [ 1, %start ]
@@ -32,7 +32,7 @@ pub unsafe fn asm_goto_with_outputs() -> u64 {
3232
pub unsafe fn asm_goto_with_outputs_use_in_label() -> u64 {
3333
let out: u64;
3434
// CHECK: [[RES:%[0-9]+]] = callbr i64 asm sideeffect alignstack inteldialect "
35-
// CHECK-NEXT: to label %[[FALLTHROUGHBB:[a-b0-9]+]] [label %[[JUMPBB:[a-b0-9]+]]]
35+
// CHECK-NEXT: to label %[[FALLTHROUGHBB:[a-b0-9]+]] [label %[[JUMPBB:return]]]
3636
asm!("{} /* {} */", out(reg) out, label { return out; });
3737
// CHECK: [[JUMPBB]]:
3838
// CHECK-NEXT: [[RET:%.+]] = phi i64 [ 1, %[[FALLTHROUGHBB]] ], [ [[RES]], %start ]
@@ -55,7 +55,7 @@ pub unsafe fn asm_goto_noreturn() -> u64 {
5555
pub unsafe fn asm_goto_noreturn_with_outputs() -> u64 {
5656
let out: u64;
5757
// CHECK: [[RES:%[0-9]+]] = callbr i64 asm sideeffect alignstack inteldialect "
58-
// CHECK-NEXT: to label %[[FALLTHROUGHBB:[a-b0-9]+]] [label %[[JUMPBB:[a-b0-9]+]]]
58+
// CHECK-NEXT: to label %[[FALLTHROUGHBB:return]] [label %[[JUMPBB:return]]]
5959
asm!("mov {}, 1", "jmp {}", out(reg) out, label { return out; });
6060
// CHECK: [[JUMPBB]]:
6161
// CHECK-NEXT: ret i64 [[RES]]

‎tests/codegen/issues/issue-112509-slice-get-andthen-get.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33

44
// CHECK-LABEL: @write_u8_variant_a
55
// CHECK-NEXT: {{.*}}:
6-
// CHECK-NEXT: icmp ugt
76
// CHECK-NEXT: getelementptr
7+
// CHECK-NEXT: icmp ugt
88
// CHECK-NEXT: select i1 {{.+}} null
99
// CHECK-NEXT: insertvalue
1010
// CHECK-NEXT: insertvalue

‎tests/codegen/multiple_returns.rs

+34
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
//@ compile-flags: -Copt-level=3 -C no-prepopulate-passes
2+
3+
#![crate_type = "lib"]
4+
5+
// CHECK-LABEL: @simple_is_one_block
6+
#[no_mangle]
7+
pub unsafe fn simple_is_one_block(x: i32) -> i32 {
8+
// CHECK: start:
9+
// CHECK-NEXT: ret i32 %x
10+
11+
// CHECK-NOT: return
12+
13+
x
14+
}
15+
16+
// CHECK-LABEL: @branch_has_shared_block
17+
#[no_mangle]
18+
pub unsafe fn branch_has_shared_block(b: bool) -> i32 {
19+
// CHECK: start:
20+
// CHECK-NEXT: %[[A:.+]] = alloca [4 x i8]
21+
// CHECK-NEXT: br i1 %b
22+
23+
// CHECK: store i32 {{42|2015}}, ptr %[[A]]
24+
// CHECK-NEXT: br label %return
25+
26+
// CHECK: store i32 {{42|2015}}, ptr %[[A]]
27+
// CHECK-NEXT: br label %return
28+
29+
// CHECK: return:
30+
// CHECK-NEXT: %[[R:.+]] = load i32, ptr %[[A]]
31+
// CHECK-NEXT: ret i32 %[[R]]
32+
33+
if b { 42 } else { 2015 }
34+
}

‎tests/mir-opt/inline/issue_106141.outer.Inline.panic-abort.diff

+15-10
Original file line numberDiff line numberDiff line change
@@ -20,30 +20,35 @@
2020
+ StorageLive(_1);
2121
+ StorageLive(_2);
2222
+ _1 = const inner::promoted[0];
23-
+ _0 = index() -> [return: bb1, unwind unreachable];
23+
+ _0 = index() -> [return: bb2, unwind unreachable];
2424
}
2525

2626
bb1: {
27+
+ StorageDead(_2);
28+
+ StorageDead(_1);
29+
return;
30+
+ }
31+
+
32+
+ bb2: {
2733
+ StorageLive(_3);
2834
+ _2 = Lt(copy _0, const 1_usize);
29-
+ assert(move _2, "index out of bounds: the length is {} but the index is {}", const 1_usize, copy _0) -> [success: bb2, unwind unreachable];
35+
+ assert(move _2, "index out of bounds: the length is {} but the index is {}", const 1_usize, copy _0) -> [success: bb3, unwind unreachable];
3036
+ }
3137
+
32-
+ bb2: {
38+
+ bb3: {
3339
+ _3 = copy (*_1)[_0];
34-
+ switchInt(move _3) -> [0: bb3, otherwise: bb4];
40+
+ switchInt(move _3) -> [0: bb4, otherwise: bb5];
3541
+ }
3642
+
37-
+ bb3: {
43+
+ bb4: {
3844
+ _0 = const 0_usize;
39-
+ goto -> bb4;
45+
+ StorageDead(_3);
46+
+ goto -> bb1;
4047
+ }
4148
+
42-
+ bb4: {
49+
+ bb5: {
4350
+ StorageDead(_3);
44-
+ StorageDead(_2);
45-
+ StorageDead(_1);
46-
return;
51+
+ goto -> bb1;
4752
}
4853
}
4954

‎tests/mir-opt/inline/issue_106141.outer.Inline.panic-unwind.diff

+15-10
Original file line numberDiff line numberDiff line change
@@ -20,30 +20,35 @@
2020
+ StorageLive(_1);
2121
+ StorageLive(_2);
2222
+ _1 = const inner::promoted[0];
23-
+ _0 = index() -> [return: bb1, unwind continue];
23+
+ _0 = index() -> [return: bb2, unwind continue];
2424
}
2525

2626
bb1: {
27+
+ StorageDead(_2);
28+
+ StorageDead(_1);
29+
return;
30+
+ }
31+
+
32+
+ bb2: {
2733
+ StorageLive(_3);
2834
+ _2 = Lt(copy _0, const 1_usize);
29-
+ assert(move _2, "index out of bounds: the length is {} but the index is {}", const 1_usize, copy _0) -> [success: bb2, unwind continue];
35+
+ assert(move _2, "index out of bounds: the length is {} but the index is {}", const 1_usize, copy _0) -> [success: bb3, unwind continue];
3036
+ }
3137
+
32-
+ bb2: {
38+
+ bb3: {
3339
+ _3 = copy (*_1)[_0];
34-
+ switchInt(move _3) -> [0: bb3, otherwise: bb4];
40+
+ switchInt(move _3) -> [0: bb4, otherwise: bb5];
3541
+ }
3642
+
37-
+ bb3: {
43+
+ bb4: {
3844
+ _0 = const 0_usize;
39-
+ goto -> bb4;
45+
+ StorageDead(_3);
46+
+ goto -> bb1;
4047
+ }
4148
+
42-
+ bb4: {
49+
+ bb5: {
4350
+ StorageDead(_3);
44-
+ StorageDead(_2);
45-
+ StorageDead(_1);
46-
return;
51+
+ goto -> bb1;
4752
}
4853
}
4954

‎tests/mir-opt/issues/issue_59352.num_to_digit.PreCodegen.after.32bit.panic-abort.mir

+1-5
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ fn num_to_digit(_1: char) -> u32 {
4949
_0 = move ((_4 as Some).0: u32);
5050
StorageDead(_5);
5151
StorageDead(_4);
52-
goto -> bb8;
52+
return;
5353
}
5454

5555
bb6: {
@@ -59,10 +59,6 @@ fn num_to_digit(_1: char) -> u32 {
5959
bb7: {
6060
StorageDead(_3);
6161
_0 = const 0_u32;
62-
goto -> bb8;
63-
}
64-
65-
bb8: {
6662
return;
6763
}
6864
}

‎tests/mir-opt/issues/issue_59352.num_to_digit.PreCodegen.after.32bit.panic-unwind.mir

+1-5
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ fn num_to_digit(_1: char) -> u32 {
4949
_0 = move ((_4 as Some).0: u32);
5050
StorageDead(_5);
5151
StorageDead(_4);
52-
goto -> bb8;
52+
return;
5353
}
5454

5555
bb6: {
@@ -59,10 +59,6 @@ fn num_to_digit(_1: char) -> u32 {
5959
bb7: {
6060
StorageDead(_3);
6161
_0 = const 0_u32;
62-
goto -> bb8;
63-
}
64-
65-
bb8: {
6662
return;
6763
}
6864
}

‎tests/mir-opt/issues/issue_59352.num_to_digit.PreCodegen.after.64bit.panic-abort.mir

+1-5
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ fn num_to_digit(_1: char) -> u32 {
4949
_0 = move ((_4 as Some).0: u32);
5050
StorageDead(_5);
5151
StorageDead(_4);
52-
goto -> bb8;
52+
return;
5353
}
5454

5555
bb6: {
@@ -59,10 +59,6 @@ fn num_to_digit(_1: char) -> u32 {
5959
bb7: {
6060
StorageDead(_3);
6161
_0 = const 0_u32;
62-
goto -> bb8;
63-
}
64-
65-
bb8: {
6662
return;
6763
}
6864
}

0 commit comments

Comments
 (0)