Skip to content

Commit 04bb561

Browse files
simonvandelsimonvandel
committed
Implement optimization that turns Eq-Not pair into Ne
1 parent 6e19e8f commit 04bb561

9 files changed

+294
-134
lines changed

‎compiler/rustc_mir/src/transform/instcombine.rs

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ use crate::transform::{MirPass, MirSource};
44
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
55
use rustc_hir::Mutability;
66
use rustc_index::vec::Idx;
7+
use rustc_middle::mir::UnOp;
78
use rustc_middle::mir::{
89
visit::PlaceContext,
910
visit::{MutVisitor, Visitor},
@@ -14,6 +15,7 @@ use rustc_middle::mir::{
1415
Rvalue,
1516
};
1617
use rustc_middle::ty::{self, TyCtxt};
18+
use smallvec::SmallVec;
1719
use std::mem;
1820

1921
pub struct InstCombine;
@@ -29,8 +31,28 @@ impl<'tcx> MirPass<'tcx> for InstCombine {
2931
optimization_finder.optimizations
3032
};
3133

34+
// Since eq_not has elements removed in the visitor, we clone it here,
35+
// such that we can still do the post visitor cleanup.
36+
let clone_eq_not = optimizations.eq_not.clone();
3237
// Then carry out those optimizations.
3338
MutVisitor::visit_body(&mut InstCombineVisitor { optimizations, tcx }, body);
39+
eq_not_post_visitor_mutations(body, clone_eq_not);
40+
}
41+
}
42+
43+
fn eq_not_post_visitor_mutations<'tcx>(
44+
body: &mut Body<'tcx>,
45+
eq_not_opts: FxHashMap<Location, EqNotOptInfo<'tcx>>,
46+
) {
47+
for (location, eq_not_opt_info) in eq_not_opts.iter() {
48+
let statements = &mut body.basic_blocks_mut()[location.block].statements;
49+
// We have to make sure that Ne is before any StorageDead as the operand being killed is used in the Ne
50+
if let Some(storage_dead_idx_to_swap_with) = eq_not_opt_info.storage_dead_to_swap_with_ne {
51+
statements.swap(location.statement_index, storage_dead_idx_to_swap_with);
52+
}
53+
if let Some(eq_stmt_idx) = eq_not_opt_info.can_remove_eq {
54+
statements[eq_stmt_idx].make_nop();
55+
}
3456
}
3557
}
3658

@@ -81,6 +103,11 @@ impl<'tcx> MutVisitor<'tcx> for InstCombineVisitor<'tcx> {
81103
*rvalue = Rvalue::Use(Operand::Copy(place));
82104
}
83105

106+
if let Some(eq_not_opt_info) = self.optimizations.eq_not.remove(&location) {
107+
*rvalue = Rvalue::BinaryOp(BinOp::Ne, eq_not_opt_info.op1, eq_not_opt_info.op2);
108+
debug!("replacing Eq and Not with {:?}", rvalue);
109+
}
110+
84111
self.super_rvalue(rvalue, location)
85112
}
86113
}
@@ -221,6 +248,85 @@ impl OptimizationFinder<'b, 'tcx> {
221248
}
222249
}
223250

251+
fn find_eq_not(&mut self, rvalue: &Rvalue<'tcx>, location: Location) -> Option<()> {
252+
// Optimize the sequence
253+
// _4 = Eq(move _5, const 2_u8);
254+
// StorageDead(_5);
255+
// _3 = Not(move _4);
256+
//
257+
// into _3 = Ne(move _5, const 2_u8)
258+
if let Rvalue::UnaryOp(UnOp::Not, op) = rvalue {
259+
let place = op.place()?;
260+
// See if we can find a Eq that assigns `place`.
261+
// We limit the search to 3 statements lookback.
262+
// Usually the first 2 statements are `StorageDead`s for operands for Eq.
263+
// We record what is marked dead so that we can reorder StorageDead so it comes after Ne
264+
265+
// We will maximum see 2 StorageDeads
266+
let mut seen_storage_deads: SmallVec<[_; 2]> = SmallVec::new();
267+
let lower_index = location.statement_index.saturating_sub(3);
268+
for (stmt_idx, stmt) in self.body.basic_blocks()[location.block].statements
269+
[lower_index..location.statement_index]
270+
.iter()
271+
.enumerate()
272+
.rev()
273+
{
274+
match &stmt.kind {
275+
rustc_middle::mir::StatementKind::Assign(box (l, r)) => {
276+
if *l == place {
277+
match r {
278+
Rvalue::BinaryOp(BinOp::Eq, op1, op2) => {
279+
// We need to make sure that the StorageDeads we saw are for
280+
// either `op1`or `op2` of Eq. Else we bail the optimization.
281+
for (dead_local, _) in seen_storage_deads.iter() {
282+
let dead_local_matches = [op1, op2].iter().any(|x| {
283+
Some(*dead_local) == x.place().map(|x| x.local)
284+
});
285+
if !dead_local_matches {
286+
return None;
287+
}
288+
}
289+
290+
// We need to make sure that the Ne we are going to insert comes before the
291+
// StorageDeads so we want to swap the StorageDead closest to Eq with Ne.
292+
let storage_dead_to_swap =
293+
seen_storage_deads.last().map(|(_, idx)| *idx);
294+
295+
// If the operand of Not is moved into it,
296+
// and that same operand is the lhs of the Eq assignment,
297+
// then we can safely remove the Eq
298+
let can_remove_eq = if op.is_move() {
299+
Some(stmt_idx + lower_index)
300+
} else {
301+
None
302+
};
303+
304+
self.optimizations.eq_not.insert(
305+
location,
306+
EqNotOptInfo {
307+
op1: op1.clone(),
308+
op2: op2.clone(),
309+
storage_dead_to_swap_with_ne: storage_dead_to_swap,
310+
can_remove_eq,
311+
},
312+
);
313+
return Some(());
314+
}
315+
_ => {}
316+
}
317+
}
318+
}
319+
rustc_middle::mir::StatementKind::StorageDead(dead) => {
320+
seen_storage_deads.push((*dead, stmt_idx + lower_index));
321+
}
322+
// If we see a pattern other than (0..=2) StorageDeads and then an Eq assignment, we conservatively bail
323+
_ => return None,
324+
}
325+
}
326+
}
327+
Some(())
328+
}
329+
224330
fn find_operand_in_equality_comparison_pattern(
225331
&self,
226332
l: &Operand<'tcx>,
@@ -265,16 +371,27 @@ impl Visitor<'tcx> for OptimizationFinder<'b, 'tcx> {
265371

266372
let _ = self.find_deref_of_address(rvalue, location);
267373

374+
let _ = self.find_eq_not(rvalue, location);
375+
268376
self.find_unneeded_equality_comparison(rvalue, location);
269377

270378
self.super_rvalue(rvalue, location)
271379
}
272380
}
273381

382+
#[derive(Clone)]
383+
struct EqNotOptInfo<'tcx> {
384+
op1: Operand<'tcx>,
385+
op2: Operand<'tcx>,
386+
storage_dead_to_swap_with_ne: Option<usize>,
387+
can_remove_eq: Option<usize>,
388+
}
389+
274390
#[derive(Default)]
275391
struct OptimizationList<'tcx> {
276392
and_stars: FxHashSet<Location>,
277393
arrays_lengths: FxHashMap<Location, Constant<'tcx>>,
278394
unneeded_equality_comparison: FxHashMap<Location, Operand<'tcx>>,
279395
unneeded_deref: FxHashMap<Location, Place<'tcx>>,
396+
eq_not: FxHashMap<Location, EqNotOptInfo<'tcx>>,
280397
}

‎src/test/mir-opt/eq_not.opt_both_moved.InstCombine.diff

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
- // MIR for `opt_both_moved` before InstCombine
2+
+ // MIR for `opt_both_moved` after InstCombine
3+
4+
fn opt_both_moved(_1: u8) -> () {
5+
debug x => _1; // in scope 0 at $DIR/eq_not.rs:18:19: 18:20
6+
let mut _0: (); // return place in scope 0 at $DIR/eq_not.rs:18:26: 18:26
7+
let _2: (); // in scope 0 at $DIR/eq_not.rs:19:5: 19:21
8+
let mut _3: bool; // in scope 0 at $DIR/eq_not.rs:19:5: 19:21
9+
let mut _4: bool; // in scope 0 at $DIR/eq_not.rs:19:13: 19:19
10+
let mut _5: u8; // in scope 0 at $DIR/eq_not.rs:19:13: 19:14
11+
let mut _6: u8; // in scope 0 at $DIR/eq_not.rs:19:18: 19:19
12+
let mut _7: !; // in scope 0 at $SRC_DIR/std/src/macros.rs:LL:COL
13+
14+
bb0: {
15+
StorageLive(_2); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
16+
StorageLive(_3); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
17+
StorageLive(_4); // scope 0 at $DIR/eq_not.rs:19:13: 19:19
18+
StorageLive(_5); // scope 0 at $DIR/eq_not.rs:19:13: 19:14
19+
_5 = _1; // scope 0 at $DIR/eq_not.rs:19:13: 19:14
20+
StorageLive(_6); // scope 0 at $DIR/eq_not.rs:19:18: 19:19
21+
_6 = _1; // scope 0 at $DIR/eq_not.rs:19:18: 19:19
22+
- _4 = Eq(move _5, move _6); // scope 0 at $DIR/eq_not.rs:19:13: 19:19
23+
- StorageDead(_6); // scope 0 at $DIR/eq_not.rs:19:18: 19:19
24+
+ nop; // scope 0 at $DIR/eq_not.rs:19:13: 19:19
25+
+ _3 = Ne(move _5, move _6); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
26+
StorageDead(_5); // scope 0 at $DIR/eq_not.rs:19:18: 19:19
27+
- _3 = Not(move _4); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
28+
+ StorageDead(_6); // scope 0 at $DIR/eq_not.rs:19:18: 19:19
29+
StorageDead(_4); // scope 0 at $DIR/eq_not.rs:19:20: 19:21
30+
switchInt(_3) -> [false: bb1, otherwise: bb2]; // scope 0 at $DIR/eq_not.rs:19:5: 19:21
31+
}
32+
33+
bb1: {
34+
_2 = const (); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
35+
StorageDead(_3); // scope 0 at $DIR/eq_not.rs:19:20: 19:21
36+
StorageDead(_2); // scope 0 at $DIR/eq_not.rs:19:20: 19:21
37+
_0 = const (); // scope 0 at $DIR/eq_not.rs:18:26: 20:2
38+
return; // scope 0 at $DIR/eq_not.rs:20:2: 20:2
39+
}
40+
41+
bb2: {
42+
StorageLive(_7); // scope 0 at $SRC_DIR/std/src/macros.rs:LL:COL
43+
begin_panic::<&str>(const "assertion failed: x == x"); // scope 0 at $SRC_DIR/std/src/macros.rs:LL:COL
44+
// mir::Constant
45+
// + span: $SRC_DIR/std/src/macros.rs:LL:COL
46+
// + literal: Const { ty: fn(&str) -> ! {std::rt::begin_panic::<&str>}, val: Value(Scalar(<ZST>)) }
47+
// ty::Const
48+
// + ty: &str
49+
// + val: Value(Slice { data: Allocation { bytes: [97, 115, 115, 101, 114, 116, 105, 111, 110, 32, 102, 97, 105, 108, 101, 100, 58, 32, 120, 32, 61, 61, 32, 120], relocations: Relocations(SortedMap { data: [] }), init_mask: InitMask { blocks: [16777215], len: Size { raw: 24 } }, size: Size { raw: 24 }, align: Align { pow2: 0 }, mutability: Not, extra: () }, start: 0, end: 24 })
50+
// mir::Constant
51+
// + span: $DIR/eq_not.rs:1:1: 1:1
52+
// + literal: Const { ty: &str, val: Value(Slice { data: Allocation { bytes: [97, 115, 115, 101, 114, 116, 105, 111, 110, 32, 102, 97, 105, 108, 101, 100, 58, 32, 120, 32, 61, 61, 32, 120], relocations: Relocations(SortedMap { data: [] }), init_mask: InitMask { blocks: [16777215], len: Size { raw: 24 } }, size: Size { raw: 24 }, align: Align { pow2: 0 }, mutability: Not, extra: () }, start: 0, end: 24 }) }
53+
}
54+
}
55+

‎src/test/mir-opt/eq_not.opt_has_later_use.InstCombine.diff

Lines changed: 33 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -4,77 +4,58 @@
44
fn opt_has_later_use(_1: Vec<u8>) -> u8 {
55
debug x => _1; // in scope 0 at $DIR/eq_not.rs:12:22: 12:23
66
let mut _0: u8; // return place in scope 0 at $DIR/eq_not.rs:12:37: 12:39
7-
let _2: (); // in scope 0 at $DIR/eq_not.rs:13:5: 13:27
8-
let mut _3: bool; // in scope 0 at $DIR/eq_not.rs:13:5: 13:27
9-
let mut _4: bool; // in scope 0 at $DIR/eq_not.rs:13:13: 13:25
10-
let mut _5: usize; // in scope 0 at $DIR/eq_not.rs:13:13: 13:20
11-
let mut _6: &std::vec::Vec<u8>; // in scope 0 at $DIR/eq_not.rs:13:13: 13:14
12-
let mut _7: !; // in scope 0 at $SRC_DIR/std/src/macros.rs:LL:COL
13-
let mut _8: &u8; // in scope 0 at $DIR/eq_not.rs:14:5: 14:9
14-
let mut _9: &std::vec::Vec<u8>; // in scope 0 at $DIR/eq_not.rs:14:5: 14:6
7+
let _2: bool; // in scope 0 at $DIR/eq_not.rs:13:9: 13:10
8+
let mut _3: bool; // in scope 0 at $DIR/eq_not.rs:13:14: 13:28
9+
let mut _4: usize; // in scope 0 at $DIR/eq_not.rs:13:15: 13:22
10+
let mut _5: &std::vec::Vec<u8>; // in scope 0 at $DIR/eq_not.rs:13:15: 13:16
11+
let mut _6: bool; // in scope 0 at $DIR/eq_not.rs:14:8: 14:9
1512
scope 1 {
16-
debug self => _6; // in scope 1 at $SRC_DIR/alloc/src/vec.rs:LL:COL
13+
debug x => _2; // in scope 1 at $DIR/eq_not.rs:13:9: 13:10
14+
}
15+
scope 2 {
16+
debug self => _5; // in scope 2 at $SRC_DIR/alloc/src/vec.rs:LL:COL
1717
}
1818

1919
bb0: {
20-
StorageLive(_2); // scope 0 at $DIR/eq_not.rs:13:5: 13:27
21-
StorageLive(_3); // scope 0 at $DIR/eq_not.rs:13:5: 13:27
22-
StorageLive(_4); // scope 0 at $DIR/eq_not.rs:13:13: 13:25
23-
StorageLive(_5); // scope 0 at $DIR/eq_not.rs:13:13: 13:20
24-
StorageLive(_6); // scope 0 at $DIR/eq_not.rs:13:13: 13:14
25-
_6 = &_1; // scope 0 at $DIR/eq_not.rs:13:13: 13:14
26-
_5 = ((*_6).1: usize); // scope 1 at $SRC_DIR/alloc/src/vec.rs:LL:COL
27-
StorageDead(_6); // scope 0 at $DIR/eq_not.rs:13:19: 13:20
28-
_4 = Eq(move _5, const 2_usize); // scope 0 at $DIR/eq_not.rs:13:13: 13:25
29-
StorageDead(_5); // scope 0 at $DIR/eq_not.rs:13:24: 13:25
30-
_3 = Not(move _4); // scope 0 at $DIR/eq_not.rs:13:5: 13:27
31-
StorageDead(_4); // scope 0 at $DIR/eq_not.rs:13:26: 13:27
32-
switchInt(_3) -> [false: bb3, otherwise: bb4]; // scope 0 at $DIR/eq_not.rs:13:5: 13:27
20+
StorageLive(_2); // scope 0 at $DIR/eq_not.rs:13:9: 13:10
21+
StorageLive(_3); // scope 0 at $DIR/eq_not.rs:13:14: 13:28
22+
StorageLive(_4); // scope 0 at $DIR/eq_not.rs:13:15: 13:22
23+
StorageLive(_5); // scope 0 at $DIR/eq_not.rs:13:15: 13:16
24+
_5 = &_1; // scope 0 at $DIR/eq_not.rs:13:15: 13:16
25+
_4 = ((*_5).1: usize); // scope 2 at $SRC_DIR/alloc/src/vec.rs:LL:COL
26+
StorageDead(_5); // scope 0 at $DIR/eq_not.rs:13:21: 13:22
27+
- _3 = Eq(move _4, const 2_usize); // scope 0 at $DIR/eq_not.rs:13:14: 13:28
28+
+ nop; // scope 0 at $DIR/eq_not.rs:13:14: 13:28
29+
+ _2 = Ne(move _4, const 2_usize); // scope 0 at $DIR/eq_not.rs:13:13: 13:28
30+
StorageDead(_4); // scope 0 at $DIR/eq_not.rs:13:27: 13:28
31+
- _2 = Not(move _3); // scope 0 at $DIR/eq_not.rs:13:13: 13:28
32+
StorageDead(_3); // scope 0 at $DIR/eq_not.rs:13:27: 13:28
33+
StorageLive(_6); // scope 1 at $DIR/eq_not.rs:14:8: 14:9
34+
_6 = _2; // scope 1 at $DIR/eq_not.rs:14:8: 14:9
35+
switchInt(_6) -> [false: bb2, otherwise: bb3]; // scope 1 at $DIR/eq_not.rs:14:5: 14:26
3336
}
3437

3538
bb1 (cleanup): {
3639
resume; // scope 0 at $DIR/eq_not.rs:12:1: 15:2
3740
}
3841

39-
bb2 (cleanup): {
40-
drop(_1) -> bb1; // scope 0 at $DIR/eq_not.rs:15:1: 15:2
42+
bb2: {
43+
_0 = const 1_u8; // scope 1 at $DIR/eq_not.rs:14:23: 14:24
44+
goto -> bb4; // scope 1 at $DIR/eq_not.rs:14:5: 14:26
4145
}
4246

4347
bb3: {
44-
_2 = const (); // scope 0 at $DIR/eq_not.rs:13:5: 13:27
45-
StorageDead(_3); // scope 0 at $DIR/eq_not.rs:13:26: 13:27
46-
StorageDead(_2); // scope 0 at $DIR/eq_not.rs:13:26: 13:27
47-
StorageLive(_8); // scope 0 at $DIR/eq_not.rs:14:5: 14:9
48-
StorageLive(_9); // scope 0 at $DIR/eq_not.rs:14:5: 14:6
49-
_9 = &_1; // scope 0 at $DIR/eq_not.rs:14:5: 14:6
50-
_8 = <Vec<u8> as Index<usize>>::index(move _9, const 0_usize) -> [return: bb5, unwind: bb2]; // scope 0 at $DIR/eq_not.rs:14:5: 14:9
51-
// mir::Constant
52-
// + span: $DIR/eq_not.rs:14:5: 14:9
53-
// + literal: Const { ty: for<'r> fn(&'r std::vec::Vec<u8>, usize) -> &'r <std::vec::Vec<u8> as std::ops::Index<usize>>::Output {<std::vec::Vec<u8> as std::ops::Index<usize>>::index}, val: Value(Scalar(<ZST>)) }
48+
_0 = const 0_u8; // scope 1 at $DIR/eq_not.rs:14:12: 14:13
49+
goto -> bb4; // scope 1 at $DIR/eq_not.rs:14:5: 14:26
5450
}
5551

5652
bb4: {
57-
StorageLive(_7); // scope 0 at $SRC_DIR/std/src/macros.rs:LL:COL
58-
begin_panic::<&str>(const "assertion failed: x.len() == 2") -> bb2; // scope 0 at $SRC_DIR/std/src/macros.rs:LL:COL
59-
// mir::Constant
60-
// + span: $SRC_DIR/std/src/macros.rs:LL:COL
61-
// + literal: Const { ty: fn(&str) -> ! {std::rt::begin_panic::<&str>}, val: Value(Scalar(<ZST>)) }
62-
// ty::Const
63-
// + ty: &str
64-
// + val: Value(Slice { data: Allocation { bytes: [97, 115, 115, 101, 114, 116, 105, 111, 110, 32, 102, 97, 105, 108, 101, 100, 58, 32, 120, 46, 108, 101, 110, 40, 41, 32, 61, 61, 32, 50], relocations: Relocations(SortedMap { data: [] }), init_mask: InitMask { blocks: [1073741823], len: Size { raw: 30 } }, size: Size { raw: 30 }, align: Align { pow2: 0 }, mutability: Not, extra: () }, start: 0, end: 30 })
65-
// mir::Constant
66-
// + span: $DIR/eq_not.rs:1:1: 1:1
67-
// + literal: Const { ty: &str, val: Value(Slice { data: Allocation { bytes: [97, 115, 115, 101, 114, 116, 105, 111, 110, 32, 102, 97, 105, 108, 101, 100, 58, 32, 120, 46, 108, 101, 110, 40, 41, 32, 61, 61, 32, 50], relocations: Relocations(SortedMap { data: [] }), init_mask: InitMask { blocks: [1073741823], len: Size { raw: 30 } }, size: Size { raw: 30 }, align: Align { pow2: 0 }, mutability: Not, extra: () }, start: 0, end: 30 }) }
53+
StorageDead(_2); // scope 0 at $DIR/eq_not.rs:15:1: 15:2
54+
StorageDead(_6); // scope 0 at $DIR/eq_not.rs:15:1: 15:2
55+
drop(_1) -> [return: bb5, unwind: bb1]; // scope 0 at $DIR/eq_not.rs:15:1: 15:2
6856
}
6957

7058
bb5: {
71-
_0 = (*_8); // scope 0 at $DIR/eq_not.rs:14:5: 14:9
72-
StorageDead(_9); // scope 0 at $DIR/eq_not.rs:14:8: 14:9
73-
StorageDead(_8); // scope 0 at $DIR/eq_not.rs:15:1: 15:2
74-
drop(_1) -> [return: bb6, unwind: bb1]; // scope 0 at $DIR/eq_not.rs:15:1: 15:2
75-
}
76-
77-
bb6: {
7859
return; // scope 0 at $DIR/eq_not.rs:15:2: 15:2
7960
}
8061
}

0 commit comments

Comments
 (0)