Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[MIR-OPT]: Optimization that turns Eq-Not pair into Ne #77031

Closed
wants to merge 5 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
138 changes: 138 additions & 0 deletions compiler/rustc_mir/src/transform/instcombine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use crate::transform::{MirPass, MirSource};
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
use rustc_hir::Mutability;
use rustc_index::vec::Idx;
use rustc_middle::mir::UnOp;
use rustc_middle::mir::{
visit::PlaceContext,
visit::{MutVisitor, Visitor},
Expand All @@ -14,6 +15,7 @@ use rustc_middle::mir::{
Rvalue,
};
use rustc_middle::ty::{self, TyCtxt};
use smallvec::SmallVec;
use std::mem;

pub struct InstCombine;
Expand All @@ -29,8 +31,30 @@ impl<'tcx> MirPass<'tcx> for InstCombine {
optimization_finder.optimizations
};

// Since eq_not has elements removed in the visitor, we clone it here,
// such that we can still do the post visitor cleanup.
let clone_eq_not = optimizations.eq_not.clone();
// Then carry out those optimizations.
MutVisitor::visit_body(&mut InstCombineVisitor { optimizations, tcx }, body);
eq_not_post_visitor_mutations(body, clone_eq_not);
}
}

fn eq_not_post_visitor_mutations<'tcx>(
body: &mut Body<'tcx>,
eq_not_opts: FxHashMap<Location, EqNotOptInfo<'tcx>>,
) {
for (location, eq_not_opt_info) in eq_not_opts.iter() {
let statements = &mut body.basic_blocks_mut()[location.block].statements;
// We have to make sure that Ne is before any StorageDead as the operand being killed is used in the Ne
if let Some(storage_dead_idx_to_swap_with) =
eq_not_opt_info.storage_dead_idx_to_swap_with_ne
{
statements.swap(location.statement_index, storage_dead_idx_to_swap_with);
}
if let Some(eq_stmt_idx) = eq_not_opt_info.can_remove_eq {
statements[eq_stmt_idx].make_nop();
}
}
}

Expand Down Expand Up @@ -81,6 +105,11 @@ impl<'tcx> MutVisitor<'tcx> for InstCombineVisitor<'tcx> {
*rvalue = Rvalue::Use(Operand::Copy(place));
}

if let Some(eq_not_opt_info) = self.optimizations.eq_not.remove(&location) {
*rvalue = Rvalue::BinaryOp(BinOp::Ne, eq_not_opt_info.op1, eq_not_opt_info.op2);
debug!("replacing Eq and Not with {:?}", rvalue);
}

self.super_rvalue(rvalue, location)
}
}
Expand Down Expand Up @@ -221,6 +250,99 @@ impl OptimizationFinder<'b, 'tcx> {
}
}

fn find_eq_not(&mut self, rvalue: &Rvalue<'tcx>, location: Location) -> Option<()> {
// Optimize the sequence
// _4 = Eq(move _5, const 2_u8);
// StorageDead(_5);
// _3 = Not(move _4);
//
// into _3 = Ne(move _5, const 2_u8)
if let Rvalue::UnaryOp(UnOp::Not, op) = rvalue {
let place = op.place()?;
// See if we can find a Eq that assigns `place`.
// We limit the search to 3 statements lookback.
// Usually the first 2 statements are `StorageDead`s for operands for Eq.
// We record what is marked dead so that we can reorder StorageDead so it comes after Ne

// We will maximum see 2 StorageDeads
let mut seen_storage_deads: SmallVec<[_; 2]> = SmallVec::new();
let lower_index = location.statement_index.saturating_sub(3);
for (stmt_idx, stmt) in self.body.basic_blocks()[location.block].statements
[lower_index..location.statement_index]
.iter()
.enumerate()
.rev()
{
match &stmt.kind {
rustc_middle::mir::StatementKind::Assign(box (l, r)) => {
if *l == place {
match r {
// FIXME(simonvandel): extend for Ne-Not pair
Rvalue::BinaryOp(BinOp::Eq, op1, op2) => {
oli-obk marked this conversation as resolved.
Show resolved Hide resolved
// We need to make sure that the StorageDeads we saw are for
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand why this is important. I'd rather think that the StorageDead for l is important, otherwise we should not remove the Neg, but instead also copy the result of the Ne operation to l.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you expand on your statement? I don't understand your point about not removing Neq (you mean Not? or Ne?).

Now that I think about it, yeah, it is perhaps unneccessary to require that the StorageDeads are for the Eq operands. Hmm do we even need to do anything special about potential StorageDeads between Eq and Not? I think not, besides doing the swap, which I think we can just do unconditionally for simplicity. We already assert that only StorageDead statements can appear between Eq and Not, and I think simply swapping a StorageDead for a complete unrelated local with Ne should not matter.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What I'm wondering is whether it is generally simpler to replace

                                    // _4 = Eq(move _5, move _6);
                                    // StorageDead(_5);
                                    // StorageDead(_6);
                                    // _3 = Not(move _4);

with

                                    // _4 = Ne(move _5, move _6);
                                    // StorageDead(_5);
                                    // StorageDead(_6);
                                    // _3 = move _4;

since then you don't have to care about any liveness or similar and leave further changes to other optimizations.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you want to get rid of the additional _3 = move 4;, maybe another peephole optimization may be better? One that just checks for _x = foo; /*arbitrary statements not touching _x or _y*/ _y = _x; and changes it to _y = foo; /*arbitrary statements not touching _x or _y*/ nop;

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah that might be simpler to do. For simplicity I'll let other passes optimize the move

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it may be beneficial to keep the current version. With your proposal, the interaction in #77031 (comment) would break. This is only one case, yeah, but it still seems worth to keep that.

What do you think?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I personally would prefer to break that interaction and fix it in a separate PR. Otherwise we need to implement both the simpler version of this PR together with the peephole opt for trivial copy propagation.

// either `op1`or `op2` of Eq. Else we bail the optimization.
for (dead_local, _) in seen_storage_deads.iter() {
let dead_local_matches = [op1, op2].iter().any(|x| {
Some(*dead_local) == x.place().map(|x| x.local)
});
if !dead_local_matches {
return None;
}
}

// Recall that we are optimizing a sequence that looks like
// this:
// _4 = Eq(move _5, move _6);
// StorageDead(_5);
// StorageDead(_6);
// _3 = Not(move _4);
//
// If we do a naive replace of Not -> Ne, we up with this:
// StorageDead(_5);
// StorageDead(_6);
// _3 = Ne(move _5, move _6);
//
// Notice that `_5` and `_6` are marked dead before being used.
// To combat this we want to swap Ne with the StorageDead
// closest to Eq, i.e `StorageDead(_5)` in this example.
let storage_dead_to_swap =
seen_storage_deads.last().map(|(_, idx)| *idx);

// If the operand of Not is moved into it,
// and that same operand is the lhs of the Eq assignment,
// then we can safely remove the Eq
let can_remove_eq = if op.is_move() {
Some(stmt_idx + lower_index)
} else {
None
};

self.optimizations.eq_not.insert(
location,
EqNotOptInfo {
op1: op1.clone(),
op2: op2.clone(),
storage_dead_idx_to_swap_with_ne: storage_dead_to_swap,
can_remove_eq,
},
);
return Some(());
}
_ => {}
}
}
}
rustc_middle::mir::StatementKind::StorageDead(dead) => {
seen_storage_deads.push((*dead, stmt_idx + lower_index));
}
// If we see a pattern other than (0..=2) StorageDeads and then an Eq assignment, we conservatively bail
_ => return None,
}
}
}
Some(())
}

fn find_operand_in_equality_comparison_pattern(
&self,
l: &Operand<'tcx>,
Expand Down Expand Up @@ -265,16 +387,32 @@ impl Visitor<'tcx> for OptimizationFinder<'b, 'tcx> {

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

let _ = self.find_eq_not(rvalue, location);

self.find_unneeded_equality_comparison(rvalue, location);

self.super_rvalue(rvalue, location)
}
}

#[derive(Clone)]
struct EqNotOptInfo<'tcx> {
/// First operand of the Eq in the Eq-Not pair
op1: Operand<'tcx>,
/// Second operand of the Eq in the Eq-Not pair
op2: Operand<'tcx>,
/// Statement index of the `StorageDead` we want to swap with Ne.
/// None if no `StorageDead` exists between Eq and Not pair)
storage_dead_idx_to_swap_with_ne: Option<usize>,
// Statement index of the Eq. None if it can not be removed
can_remove_eq: Option<usize>,
}

#[derive(Default)]
struct OptimizationList<'tcx> {
and_stars: FxHashSet<Location>,
arrays_lengths: FxHashMap<Location, Constant<'tcx>>,
unneeded_equality_comparison: FxHashMap<Location, Operand<'tcx>>,
unneeded_deref: FxHashMap<Location, Place<'tcx>>,
eq_not: FxHashMap<Location, EqNotOptInfo<'tcx>>,
}
55 changes: 55 additions & 0 deletions src/test/mir-opt/eq_not.opt_both_moved.InstCombine.diff
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
- // MIR for `opt_both_moved` before InstCombine
+ // MIR for `opt_both_moved` after InstCombine

fn opt_both_moved(_1: u8) -> () {
debug x => _1; // in scope 0 at $DIR/eq_not.rs:18:19: 18:20
let mut _0: (); // return place in scope 0 at $DIR/eq_not.rs:18:26: 18:26
let _2: (); // in scope 0 at $DIR/eq_not.rs:19:5: 19:21
let mut _3: bool; // in scope 0 at $DIR/eq_not.rs:19:5: 19:21
let mut _4: bool; // in scope 0 at $DIR/eq_not.rs:19:13: 19:19
let mut _5: u8; // in scope 0 at $DIR/eq_not.rs:19:13: 19:14
let mut _6: u8; // in scope 0 at $DIR/eq_not.rs:19:18: 19:19
let mut _7: !; // in scope 0 at $SRC_DIR/std/src/macros.rs:LL:COL

bb0: {
StorageLive(_2); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
StorageLive(_3); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
StorageLive(_4); // scope 0 at $DIR/eq_not.rs:19:13: 19:19
StorageLive(_5); // scope 0 at $DIR/eq_not.rs:19:13: 19:14
_5 = _1; // scope 0 at $DIR/eq_not.rs:19:13: 19:14
StorageLive(_6); // scope 0 at $DIR/eq_not.rs:19:18: 19:19
_6 = _1; // scope 0 at $DIR/eq_not.rs:19:18: 19:19
- _4 = Eq(move _5, move _6); // scope 0 at $DIR/eq_not.rs:19:13: 19:19
- StorageDead(_6); // scope 0 at $DIR/eq_not.rs:19:18: 19:19
+ nop; // scope 0 at $DIR/eq_not.rs:19:13: 19:19
+ _3 = Ne(move _5, move _6); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
StorageDead(_5); // scope 0 at $DIR/eq_not.rs:19:18: 19:19
- _3 = Not(move _4); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
+ StorageDead(_6); // scope 0 at $DIR/eq_not.rs:19:18: 19:19
StorageDead(_4); // scope 0 at $DIR/eq_not.rs:19:20: 19:21
switchInt(_3) -> [false: bb1, otherwise: bb2]; // scope 0 at $DIR/eq_not.rs:19:5: 19:21
}

bb1: {
_2 = const (); // scope 0 at $DIR/eq_not.rs:19:5: 19:21
StorageDead(_3); // scope 0 at $DIR/eq_not.rs:19:20: 19:21
StorageDead(_2); // scope 0 at $DIR/eq_not.rs:19:20: 19:21
_0 = const (); // scope 0 at $DIR/eq_not.rs:18:26: 20:2
return; // scope 0 at $DIR/eq_not.rs:20:2: 20:2
}

bb2: {
StorageLive(_7); // scope 0 at $SRC_DIR/std/src/macros.rs:LL:COL
begin_panic::<&str>(const "assertion failed: x == x"); // scope 0 at $SRC_DIR/std/src/macros.rs:LL:COL
// mir::Constant
// + span: $SRC_DIR/std/src/macros.rs:LL:COL
// + literal: Const { ty: fn(&str) -> ! {std::rt::begin_panic::<&str>}, val: Value(Scalar(<ZST>)) }
// ty::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 })
// mir::Constant
// + span: $DIR/eq_not.rs:1:1: 1:1
// + 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 }) }
}
}

62 changes: 62 additions & 0 deletions src/test/mir-opt/eq_not.opt_has_later_use.InstCombine.diff
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
- // MIR for `opt_has_later_use` before InstCombine
+ // MIR for `opt_has_later_use` after InstCombine

fn opt_has_later_use(_1: Vec<u8>) -> u8 {
debug x => _1; // in scope 0 at $DIR/eq_not.rs:12:22: 12:23
let mut _0: u8; // return place in scope 0 at $DIR/eq_not.rs:12:37: 12:39
let _2: bool; // in scope 0 at $DIR/eq_not.rs:13:9: 13:10
let mut _3: bool; // in scope 0 at $DIR/eq_not.rs:13:14: 13:28
let mut _4: usize; // in scope 0 at $DIR/eq_not.rs:13:15: 13:22
let mut _5: &std::vec::Vec<u8>; // in scope 0 at $DIR/eq_not.rs:13:15: 13:16
let mut _6: bool; // in scope 0 at $DIR/eq_not.rs:14:8: 14:9
scope 1 {
debug x => _2; // in scope 1 at $DIR/eq_not.rs:13:9: 13:10
}
scope 2 {
debug self => _5; // in scope 2 at $SRC_DIR/alloc/src/vec.rs:LL:COL
}

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

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

bb2: {
_0 = const 1_u8; // scope 1 at $DIR/eq_not.rs:14:23: 14:24
goto -> bb4; // scope 1 at $DIR/eq_not.rs:14:5: 14:26
}

bb3: {
_0 = const 0_u8; // scope 1 at $DIR/eq_not.rs:14:12: 14:13
goto -> bb4; // scope 1 at $DIR/eq_not.rs:14:5: 14:26
}

bb4: {
StorageDead(_2); // scope 0 at $DIR/eq_not.rs:15:1: 15:2
StorageDead(_6); // scope 0 at $DIR/eq_not.rs:15:1: 15:2
drop(_1) -> [return: bb5, unwind: bb1]; // scope 0 at $DIR/eq_not.rs:15:1: 15:2
}

bb5: {
return; // scope 0 at $DIR/eq_not.rs:15:2: 15:2
}
}

Loading