Skip to content

Commit be9306c

Browse files
borsbors
committed
Auto merge of #138144 - scottmcm:multiple-mir-returns, r=<try>
Use multiple returns in MIR if it saves a block; still have only one in LLVM This is still an open question whether it's desired, per #72022 (comment), so opening as a draft. Given that cranelift prefers multiple returns and we use optimized MIR for MIR inlining, it seems at least plausible that MIR shouldn't have the "one return" rule (the validator doesn't enforce it currently either) so we can have fewer total blocks, and then targets like LLVM can enforce it if they so choose (as this PR does). r? ghost
2 parents 91a0e16 + 2dd720c commit be9306c

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)