Skip to content

Commit d24d035

Browse files
ohadravidohadravid
committed
Use MaybeStorageDead analysis in copy_prop mir-opt to remove fewer storage statements
1 parent 375eda2 commit d24d035

21 files changed

+203
-35
lines changed

‎compiler/rustc_mir_transform/src/copy_prop.rs

Lines changed: 59 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
1+
use std::borrow::Cow;
2+
13
use rustc_index::IndexSlice;
24
use rustc_index::bit_set::DenseBitSet;
35
use rustc_middle::mir::visit::*;
46
use rustc_middle::mir::*;
57
use rustc_middle::ty::TyCtxt;
8+
use rustc_mir_dataflow::impls::{MaybeStorageDead, always_storage_live_locals};
9+
use rustc_mir_dataflow::{Analysis, ResultsCursor};
610
use tracing::{debug, instrument};
711

812
use crate::ssa::SsaLocals;
@@ -16,7 +20,7 @@ use crate::ssa::SsaLocals;
1620
/// _d = move? _c
1721
/// where each of the locals is only assigned once.
1822
///
19-
/// We want to replace all those locals by `_a`, either copied or moved.
23+
/// We want to replace all those locals by `_a` (the "head"), either copied or moved.
2024
pub(super) struct CopyProp;
2125

2226
impl<'tcx> crate::MirPass<'tcx> for CopyProp {
@@ -34,15 +38,40 @@ impl<'tcx> crate::MirPass<'tcx> for CopyProp {
3438
let fully_moved = fully_moved_locals(&ssa, body);
3539
debug!(?fully_moved);
3640

37-
let mut storage_to_remove = DenseBitSet::new_empty(fully_moved.domain_size());
41+
let mut head_storage_to_check = DenseBitSet::new_empty(fully_moved.domain_size());
42+
3843
for (local, &head) in ssa.copy_classes().iter_enumerated() {
3944
if local != head {
40-
storage_to_remove.insert(head);
45+
// We need to determine if we can keep the head's storage statements (which enables better optimizations).
46+
// For every local's usage location, if the head is in `maybe_storage_dead`, we have to remove the storage statements for it.
47+
head_storage_to_check.insert(head);
4148
}
4249
}
4350

4451
let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
4552

53+
let storage_to_remove = if any_replacement {
54+
let always_live_locals = &always_storage_live_locals(body);
55+
56+
let maybe_storage_dead = MaybeStorageDead::new(Cow::Borrowed(always_live_locals))
57+
.iterate_to_fixpoint(tcx, body, None)
58+
.into_results_cursor(body);
59+
60+
let mut storage_checker = StorageChecker {
61+
copy_classes: ssa.copy_classes(),
62+
maybe_storage_dead,
63+
head_storage_to_check,
64+
storage_to_remove: DenseBitSet::new_empty(fully_moved.domain_size()),
65+
};
66+
67+
storage_checker.visit_body(body);
68+
69+
storage_checker.storage_to_remove
70+
} else {
71+
// Will be empty anyway.
72+
head_storage_to_check
73+
};
74+
4675
Replacer {
4776
tcx,
4877
copy_classes: ssa.copy_classes(),
@@ -119,6 +148,7 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
119148
if self.borrowed_locals.contains(*local) {
120149
return;
121150
}
151+
122152
match ctxt {
123153
// Do not modify the local in storage statements.
124154
PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {}
@@ -172,3 +202,29 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
172202
}
173203
}
174204
}
205+
206+
struct StorageChecker<'a, 'tcx> {
207+
storage_to_remove: DenseBitSet<Local>,
208+
head_storage_to_check: DenseBitSet<Local>,
209+
maybe_storage_dead: ResultsCursor<'a, 'tcx, MaybeStorageDead<'a>>,
210+
copy_classes: &'a IndexSlice<Local, Local>,
211+
}
212+
213+
impl<'a, 'tcx> Visitor<'tcx> for StorageChecker<'a, 'tcx> {
214+
fn visit_local(&mut self, local: Local, context: PlaceContext, location: Location) {
215+
let head = self.copy_classes[local];
216+
217+
if context.is_use() && self.head_storage_to_check.contains(head) {
218+
self.maybe_storage_dead.seek_after_primary_effect(location);
219+
if self.maybe_storage_dead.get().contains(head) {
220+
debug!(
221+
?location,
222+
?local,
223+
?head,
224+
"found use of local with head at a location in which head is maybe dead, marking head for storage removal"
225+
);
226+
self.storage_to_remove.insert(head);
227+
}
228+
}
229+
}
230+
}

‎tests/mir-opt/copy-prop/cycle.main.CopyProp.panic-abort.diff

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
}
2727

2828
bb1: {
29-
- StorageLive(_2);
29+
StorageLive(_2);
3030
_2 = copy _1;
3131
- StorageLive(_3);
3232
- _3 = copy _2;
@@ -46,7 +46,7 @@
4646
StorageDead(_5);
4747
_0 = const ();
4848
- StorageDead(_3);
49-
- StorageDead(_2);
49+
StorageDead(_2);
5050
StorageDead(_1);
5151
return;
5252
}

‎tests/mir-opt/copy-prop/cycle.main.CopyProp.panic-unwind.diff

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
}
2727

2828
bb1: {
29-
- StorageLive(_2);
29+
StorageLive(_2);
3030
_2 = copy _1;
3131
- StorageLive(_3);
3232
- _3 = copy _2;
@@ -46,7 +46,7 @@
4646
StorageDead(_5);
4747
_0 = const ();
4848
- StorageDead(_3);
49-
- StorageDead(_2);
49+
StorageDead(_2);
5050
StorageDead(_1);
5151
return;
5252
}

‎tests/mir-opt/copy-prop/dead_stores_79191.f.CopyProp.after.panic-abort.mir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
1111
}
1212

1313
bb0: {
14+
StorageLive(_2);
1415
_2 = copy _1;
1516
_1 = const 5_usize;
1617
_1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
2122

2223
bb1: {
2324
StorageDead(_4);
25+
StorageDead(_2);
2426
return;
2527
}
2628
}

‎tests/mir-opt/copy-prop/dead_stores_79191.f.CopyProp.after.panic-unwind.mir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
1111
}
1212

1313
bb0: {
14+
StorageLive(_2);
1415
_2 = copy _1;
1516
_1 = const 5_usize;
1617
_1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
2122

2223
bb1: {
2324
StorageDead(_4);
25+
StorageDead(_2);
2426
return;
2527
}
2628
}

‎tests/mir-opt/copy-prop/dead_stores_better.f.CopyProp.after.panic-abort.mir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
1111
}
1212

1313
bb0: {
14+
StorageLive(_2);
1415
_2 = copy _1;
1516
_1 = const 5_usize;
1617
_1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
2122

2223
bb1: {
2324
StorageDead(_4);
25+
StorageDead(_2);
2426
return;
2527
}
2628
}

‎tests/mir-opt/copy-prop/dead_stores_better.f.CopyProp.after.panic-unwind.mir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ fn f(_1: usize) -> usize {
1111
}
1212

1313
bb0: {
14+
StorageLive(_2);
1415
_2 = copy _1;
1516
_1 = const 5_usize;
1617
_1 = copy _2;
@@ -21,6 +22,7 @@ fn f(_1: usize) -> usize {
2122

2223
bb1: {
2324
StorageDead(_4);
25+
StorageDead(_2);
2426
return;
2527
}
2628
}

‎tests/mir-opt/copy-prop/issue_107511.main.CopyProp.panic-abort.diff

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@
8686
}
8787

8888
bb6: {
89-
- StorageLive(_16);
89+
StorageLive(_16);
9090
_16 = copy ((_11 as Some).0: usize);
9191
StorageLive(_17);
9292
- StorageLive(_18);
@@ -116,7 +116,7 @@
116116
StorageDead(_17);
117117
- StorageDead(_18);
118118
- _10 = const ();
119-
- StorageDead(_16);
119+
StorageDead(_16);
120120
StorageDead(_13);
121121
StorageDead(_11);
122122
- StorageDead(_10);

‎tests/mir-opt/copy-prop/issue_107511.main.CopyProp.panic-unwind.diff

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@
8686
}
8787

8888
bb6: {
89-
- StorageLive(_16);
89+
StorageLive(_16);
9090
_16 = copy ((_11 as Some).0: usize);
9191
StorageLive(_17);
9292
- StorageLive(_18);
@@ -116,7 +116,7 @@
116116
StorageDead(_17);
117117
- StorageDead(_18);
118118
- _10 = const ();
119-
- StorageDead(_16);
119+
StorageDead(_16);
120120
StorageDead(_13);
121121
StorageDead(_11);
122122
- StorageDead(_10);

‎tests/mir-opt/copy-prop/issue_141649.main.CopyProp.panic-abort.diff

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
- // MIR for `main` before CopyProp
2+
+ // MIR for `main` after CopyProp
3+
4+
fn main() -> () {
5+
let mut _0: ();
6+
let _1: ();
7+
let _2: main::S;
8+
let _3: ();
9+
let mut _4: main::S;
10+
let _5: ();
11+
let _6: main::S;
12+
let _7: ();
13+
let mut _8: main::S;
14+
let _9: ();
15+
let _10: main::C;
16+
let _11: ();
17+
let mut _12: main::C;
18+
let _13: main::C;
19+
let _14: ();
20+
let mut _15: main::C;
21+
scope 1 {
22+
debug s1 => _2;
23+
}
24+
scope 2 {
25+
debug s2 => _6;
26+
}
27+
scope 3 {
28+
debug c1 => _10;
29+
}
30+
scope 4 {
31+
debug c2 => _13;
32+
}
33+
34+
bb0: {
35+
- StorageLive(_1);
36+
StorageLive(_2);
37+
_2 = S(const 1_usize, const 2_usize);
38+
StorageLive(_3);
39+
- StorageLive(_4);
40+
- _4 = move _2;
41+
- _3 = std::mem::drop::<S>(move _4) -> [return: bb1, unwind unreachable];
42+
+ _3 = std::mem::drop::<S>(move _2) -> [return: bb1, unwind unreachable];
43+
}
44+
45+
bb1: {
46+
- StorageDead(_4);
47+
StorageDead(_3);
48+
- _1 = const ();
49+
StorageDead(_2);
50+
- StorageDead(_1);
51+
- StorageLive(_5);
52+
StorageLive(_6);
53+
_6 = S(const 3_usize, const 4_usize);
54+
StorageLive(_7);
55+
- StorageLive(_8);
56+
- _8 = move _6;
57+
- _7 = std::mem::drop::<S>(move _8) -> [return: bb2, unwind unreachable];
58+
+ _7 = std::mem::drop::<S>(move _6) -> [return: bb2, unwind unreachable];
59+
}
60+
61+
bb2: {
62+
- StorageDead(_8);
63+
StorageDead(_7);
64+
- _5 = const ();
65+
StorageDead(_6);
66+
- StorageDead(_5);
67+
- StorageLive(_9);
68+
StorageLive(_10);
69+
_10 = C(const 1_usize, const 2_usize);
70+
StorageLive(_11);
71+
- StorageLive(_12);
72+
- _12 = copy _10;
73+
- _11 = std::mem::drop::<C>(move _12) -> [return: bb3, unwind unreachable];
74+
+ _11 = std::mem::drop::<C>(copy _10) -> [return: bb3, unwind unreachable];
75+
}
76+
77+
bb3: {
78+
- StorageDead(_12);
79+
StorageDead(_11);
80+
- _9 = const ();
81+
StorageDead(_10);
82+
- StorageDead(_9);
83+
StorageLive(_13);
84+
_13 = C(const 3_usize, const 4_usize);
85+
StorageLive(_14);
86+
- StorageLive(_15);
87+
- _15 = copy _13;
88+
- _14 = std::mem::drop::<C>(move _15) -> [return: bb4, unwind unreachable];
89+
+ _14 = std::mem::drop::<C>(copy _13) -> [return: bb4, unwind unreachable];
90+
}
91+
92+
bb4: {
93+
- StorageDead(_15);
94+
StorageDead(_14);
95+
_0 = const ();
96+
StorageDead(_13);
97+
return;
98+
}
99+
}
100+

‎tests/mir-opt/copy-prop/issue_141649.main.CopyProp.panic-unwind.diff

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333

3434
bb0: {
3535
- StorageLive(_1);
36-
- StorageLive(_2);
36+
StorageLive(_2);
3737
_2 = S(const 1_usize, const 2_usize);
3838
StorageLive(_3);
3939
- StorageLive(_4);
@@ -46,10 +46,10 @@
4646
- StorageDead(_4);
4747
StorageDead(_3);
4848
- _1 = const ();
49-
- StorageDead(_2);
49+
StorageDead(_2);
5050
- StorageDead(_1);
5151
- StorageLive(_5);
52-
- StorageLive(_6);
52+
StorageLive(_6);
5353
_6 = S(const 3_usize, const 4_usize);
5454
StorageLive(_7);
5555
- StorageLive(_8);
@@ -62,10 +62,10 @@
6262
- StorageDead(_8);
6363
StorageDead(_7);
6464
- _5 = const ();
65-
- StorageDead(_6);
65+
StorageDead(_6);
6666
- StorageDead(_5);
6767
- StorageLive(_9);
68-
- StorageLive(_10);
68+
StorageLive(_10);
6969
_10 = C(const 1_usize, const 2_usize);
7070
StorageLive(_11);
7171
- StorageLive(_12);
@@ -78,9 +78,9 @@
7878
- StorageDead(_12);
7979
StorageDead(_11);
8080
- _9 = const ();
81-
- StorageDead(_10);
81+
StorageDead(_10);
8282
- StorageDead(_9);
83-
- StorageLive(_13);
83+
StorageLive(_13);
8484
_13 = C(const 3_usize, const 4_usize);
8585
StorageLive(_14);
8686
- StorageLive(_15);
@@ -93,7 +93,7 @@
9393
- StorageDead(_15);
9494
StorageDead(_14);
9595
_0 = const ();
96-
- StorageDead(_13);
96+
StorageDead(_13);
9797
return;
9898
}
9999
}

0 commit comments

Comments
 (0)