Simplify switches on a statically known discriminant in MIR. #112688

JohnBobbo96 · 2023-06-16T02:23:11Z

This PR adds a new MIR pass: SimplifyStaticSwitch, which aims to simplify switches in MIR that match on a known discriminant, this mainly happens as a result of MIR inlining.

Addresses: #111442 (comment)

rustbot · 2023-06-16T02:23:22Z

Some changes occurred to MIR optimizations

cc @rust-lang/wg-mir-opt

bors · 2023-06-16T05:45:03Z

☔ The latest upstream changes (presumably #112346) made this pull request unmergeable. Please resolve the merge conflicts.

scottmcm · 2023-06-16T05:50:02Z

cc @cjgillot specifically, to comment how this should go along with #107009

cjgillot

We already have SeparateConstSwitch and ConstProp, which taken together do a similar transformation. What is the gain of this pass compared to the current state?

cjgillot · 2023-06-16T18:18:46Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        debug!(body = %tcx.def_path_str(body.source.def_id()));


Suggested change

debug!(body = %tcx.def_path_str(body.source.def_id()));

debug!(body = ?body.source.def_id());

def_path_str dooms the current compilation session.

cjgillot · 2023-06-16T18:19:01Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+#[instrument(level = "debug", skip(tcx, body), ret)]
+fn simplify_static_switch<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool {
+    let mut new_targets = FxHashMap::default();
+    let predecessors: &IndexSlice<_, _> = body.basic_blocks.predecessors();


Suggested change

let predecessors: &IndexSlice<_, _> = body.basic_blocks.predecessors();

let predecessors = body.basic_blocks.predecessors();

cjgillot · 2023-06-16T18:19:15Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+            continue;
+        }
+
+        let predecessors: &[BasicBlock] = &predecessors[block];


Suggested change

let predecessors: &[BasicBlock] = &predecessors[block];

let predecessors = &predecessors[block];

cjgillot · 2023-06-16T18:19:44Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+fn simplify_static_switch<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool {
+    let mut new_targets = FxHashMap::default();
+    let predecessors: &IndexSlice<_, _> = body.basic_blocks.predecessors();
+    for (block, data) in traversal::preorder(body) {


Why this particular order?

I thought it might be better than iterating normally but i have switched it back to just iter_enumerated.

cjgillot · 2023-06-16T18:21:14Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+                                "{pred:?}: {place:?} = {}::{variant:?}; goto -> {block:?};",
+                                tcx.def_path_str(def_id)


Suggested change

"{pred:?}: {place:?} = {}::{variant:?}; goto -> {block:?};",

tcx.def_path_str(def_id)

"{pred:?}: {place:?} = {def_id:?}::{variant:?}; goto -> {block:?};"

compiler/rustc_mir_transform/src/simplify_static_switch.rs

cjgillot · 2023-06-16T18:24:16Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+                                "{pred:?}: {place:?} = {}::{variant:?}; goto -> {block:?};",
+                                tcx.def_path_str(def_id)
+                            );
+                            let discr_ty = tcx.type_of(def_id).skip_binder();


You must take the substitutions into account here, otherwise you'll get the wrong type.

JohnBobbo96 · 2023-06-17T03:55:14Z

@rustbot author

JohnBobbo96 · 2023-06-19T22:43:51Z

We already have SeparateConstSwitch and ConstProp, which taken together do a similar transformation. What is the gain of this pass compared to the current state?

I'd say the main advantage is that this pass sort of combines the effect of SeparateConstSwitch + ConstProp except it looks for aggregates instead of constants, which currently doesn't seem to be handled by SeparateConstSwitch + ConstProp alone, see this godbolt link (the blocks that switch on the discriminant of _2 are still there, which this pass removes).

JohnBobbo96 · 2023-06-19T22:45:46Z

@rustbot ready

cjgillot · 2023-06-20T17:34:40Z

compiler/rustc_mir_transform/src/lib.rs

@@ -130,7 +131,7 @@ pub fn provide(providers: &mut Providers) {
        mir_generator_witnesses: generator::mir_generator_witnesses,
        optimized_mir,
        is_mir_available,
-        is_ctfe_mir_available: |tcx, did| is_mir_available(tcx, did),
+        is_ctfe_mir_available: is_mir_available,


This change is unrelated. Could you submit it as a separate PR?

i removed this for now, i'll open a PR soon.

cjgillot · 2023-06-20T17:34:58Z

compiler/rustc_mir_transform/src/lib.rs

@@ -532,7 +533,7 @@ fn run_runtime_cleanup_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
 }

 fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-    fn o1<T>(x: T) -> WithMinOptLevel<T> {
+    const fn o1<T>(x: T) -> WithMinOptLevel<T> {


Why is this useful?

i also removed this.

cjgillot · 2023-06-20T17:35:49Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+    fn name(&self) -> &'static str {
+        "SimplifyStaticSwitch"
+    }
+


Suggested change

fn name(&self) -> &'static str {

"SimplifyStaticSwitch"

}

This is the default.

compiler/rustc_mir_transform/src/simplify_static_switch.rs

cjgillot · 2023-06-20T17:49:23Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+    }
+
+    let basic_blocks = body.basic_blocks.as_mut();
+    for ((place, block), new_targets) in new_targets {


This is what you call discr earlier, could you name it so here?

cjgillot · 2023-06-20T17:51:42Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+            vec![]
+        };
+
+        for (block, target) in new_targets {


block here is pred, please call it so.

cjgillot · 2023-06-20T17:52:05Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+
+        for (block, target) in new_targets {
+            let data = &mut basic_blocks[block];
+            let old_goto = &mut data.terminator.as_mut().unwrap().kind;


Suggested change

let old_goto = &mut data.terminator.as_mut().unwrap().kind;

let old_goto = &mut data.terminator_mut().kind;

cjgillot · 2023-06-20T17:53:20Z

compiler/rustc_mir_transform/src/lib.rs

@@ -548,6 +549,8 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
            &uninhabited_enum_branching::UninhabitedEnumBranching,
            &o1(simplify::SimplifyCfg::AfterUninhabitedEnumBranching),
            &inline::Inline,
+            // Remove switches on a statically known discriminant, which can happen as a result of inlining.
+            &simplify_static_switch::SimplifyStaticSwitch,


This pass clones statements. As we have SSA-based passes, and we don't try to recover SSAness, it should be run much later.

i have moved the pass to right before SeparateConstSwitch.

cjgillot · 2023-06-20T17:54:55Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+                    if place.as_local() == Some(switched)
+                        && let Some(local) = discr.as_local()
+                        && !borrowed_locals.contains(local) =>


Those checks should happen outside of the find_map, in cases there are several assignments to switched.

i also moved these checks to outside of the find_map.

JohnBobbo96 · 2023-06-20T22:14:04Z

@rustbot author

JohnBobbo96 · 2023-06-20T23:34:00Z

@rustbot ready

cjgillot

Thanks! This is starting to look good.
Could you add comments explaining the pattern you are looking for in this opt, and why it is sound?

cjgillot · 2023-06-23T07:10:36Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+        // We need to make sure that `switched` is not overwritten before the switch,
+        // and also that `discr` is not overwritten before the `Rvalue::Discriminant`.
+        if is_local_mutated(body, block, discr, index, 0..index)
+            || is_local_mutated(body, block, switched, index, index..data.statements.len())


Nit: can we just pass a range to this function?

cjgillot · 2023-06-23T07:12:17Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+    let predecessors = body.basic_blocks.predecessors();
+    let borrowed_locals = borrowed_locals(body);
+    // Should this be an `IndexVec` instead?
+    let mut new_targets = FxHashMap::default();


A simple Vec should be enough, as you don't use it for indexing.

cjgillot · 2023-06-23T07:15:20Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+                    StatementKind::Assign(box (
+                        place,
+                        Rvalue::Aggregate(box AggregateKind::Adt(def_id, variant, substs, ..), ..),
+                    )) if place.as_local() == Some(discr) => {


You need to check if discr is mutated between this assignment and the end of block.

I have added a check for this.

cjgillot · 2023-06-23T07:15:33Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+            };
+
+            for statement in data.statements.iter().rev() {
+                match statement.kind {


Nit: this can be an if-let.

cjgillot · 2023-06-23T07:19:33Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+    for (block, new_targets) in new_targets {
+        let statements: Vec<_> = basic_blocks[block].statements.iter().cloned().collect();
+        for (pred, target) in new_targets {
+            let pred = &mut basic_blocks[pred];


Please document why there is only a single pred -> block edge in the whole CFG, and that no earlier threading has overwritten it.

It may also be useful to document why pred != block and assert it.

cjgillot · 2023-06-23T07:20:06Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+        let statements: Vec<_> = basic_blocks[block].statements.iter().cloned().collect();
+        for (pred, target) in new_targets {
+            let pred = &mut basic_blocks[pred];
+            let old_goto = &mut pred.terminator_mut().kind;


We should assert that old_goto is indeed goto block, to ensure we have not overwritten it in the mean time.

I have added an assert to check this, as well as not allowing two StaticSwitchs to have the same pred, so the terminator should not get overwritten.

cjgillot · 2023-06-23T07:21:26Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+
+    let basic_blocks = body.basic_blocks.as_mut();
+    for (block, new_targets) in new_targets {
+        let statements: Vec<_> = basic_blocks[block].statements.iter().cloned().collect();


Could we use IndexSlice::pick2_mut to avoid having to clone statements twice?

This is cool, i also added this.

cjgillot · 2023-06-23T07:23:55Z

tests/mir-opt/simplify_static_switch.rs

@@ -0,0 +1,29 @@
+// unit-test: SimplifyStaticSwitch


Could you add more tests?
With all the extra checks I made you add + weird cfgs if possible.
See #107009 for a few examples.

I have added a few more tests, including checks for mutation/borrowing.

cjgillot · 2023-06-23T07:24:34Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+
+            // Something to be noted is that, this creates an edge from `pred -> target`
+            // which will only appear once in the CFG.
+            *old_goto = new_goto;


Should we prevent threading through a loop header, in order to avoid creating irreducible CFGs?

I also added a check for this.

JohnBobbo96 · 2023-06-24T01:05:13Z

@rustbot author

JohnBobbo96 · 2023-06-27T04:23:36Z

Thanks! This is starting to look good. Could you add comments explaining the pattern you are looking for in this opt, and why it is sound?

Thank you for reviewing this!
I have also added more to the documentation explaining why i believe it is sound.

JohnBobbo96 · 2023-06-27T04:27:11Z

@rustbot ready

removes unnecessary switches on statically known discriminants.

cjgillot · 2023-07-04T18:51:29Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+            .statements
+            .iter()
+            .enumerate()
+            .take_while(|&(index, _)| index != statement_index)


Suggested change

.statements

.iter()

.enumerate()

.take_while(|&(index, _)| index != statement_index)

.statements[..statement_index]

.iter()

.enumerate()

cjgillot · 2023-07-04T19:03:30Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+            continue
+        };
+
+        if place.local != switched {


local has necessarily an integer type. Should we assert that place.projection is empty?

cjgillot · 2023-07-04T19:06:22Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+            continue;
+        }
+
+        let mut finder = MutatedLocalFinder { local: discr, mutated: false };


This should be created inside the loop, to ensure we don't forget to reset mutated.

cjgillot · 2023-07-04T19:07:10Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+                    | StatementKind::Assign(box (
+                        place,
+                        Rvalue::Aggregate(box AggregateKind::Adt(_, variant, ..), ..),
+                    )) if place.local == discr => {


Assert place.projection is empty here too.

cjgillot · 2023-07-04T19:08:45Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+
+                        continue 'preds;
+                    }
+                    _ => finder.visit_statement(statement, Location { block, statement_index }),


Could we create finder, visit, and check mutated in this arm, instead of the other arms?

cjgillot · 2023-07-04T19:19:20Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+                    Some(exclude) => {
+                        pred.statements.extend(block.statements.iter().enumerate().filter_map(
+                            |(index, statement)| {
+                                if index == exclude { None } else { Some(statement.clone()) }


That's not necessary. We can run remove_unused_definitions together with remove_dead_blocks.

cjgillot · 2023-07-04T19:22:48Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+    fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _: Location) {
+        if self.local == place.local && let PlaceContext::MutatingUse(..) = context {
+            self.mutated = true;
+        }


Should we detect storage statements too?

cjgillot · 2023-07-04T19:25:20Z

tests/mir-opt/separate_const_switch.identity.SeparateConstSwitch.diff

@@ -51,51 +51,44 @@
          StorageLive(_10);
          StorageLive(_11);
          _9 = discriminant(_1);
-          switchInt(move _9) -> [0: bb7, 1: bb5, otherwise: bb6];
+          switchInt(move _9) -> [0: bb5, 1: bb3, otherwise: bb4];


Should we make this a unit test for SeparateConstSwitch?

cjgillot · 2023-07-04T19:35:45Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+
+            // We use the SSA, to destroy the SSA.
+            let data = {
+                let (block, pred) = basic_blocks.pick2_mut(block, switch.pred);


Could you add a comment explaining why block != switch.pred?

cjgillot · 2023-07-04T19:36:53Z

compiler/rustc_mir_transform/src/simplify_static_switch.rs

+    }
+
+    let basic_blocks = body.basic_blocks.as_mut();
+    let num_switches: usize = static_switches.iter().map(|(_, switches)| switches.len()).sum();


Counting the number of opts just for tracing is not useful.

bors · 2023-07-14T19:38:04Z

☔ The latest upstream changes (presumably #109025) made this pull request unmergeable. Please resolve the merge conflicts.

JohnCSimon · 2023-08-27T05:33:28Z

@JohnBobbo96

ping from triage - can you post your status on this PR? There hasn't been an update in a few months and there are merge conflicts Thanks!

FYI: when a PR is ready for review, send a message containing
@rustbot ready to switch to S-waiting-on-review so the PR is in the reviewer's backlog.

JohnCSimon · 2023-12-17T22:08:50Z

@JohnBobbo96

Ping from triage: I'm closing this due to inactivity, Please reopen when you are ready to continue with this.
Note: if you are going to continue please open the PR BEFORE you push to it, else you won't be able to reopen - this is a quirk of github.
Thanks for your contribution.

@rustbot label: +S-inactive

rustbot assigned scottmcm Jun 16, 2023

rustbot added S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jun 16, 2023

cjgillot reviewed Jun 16, 2023

View reviewed changes

scottmcm assigned cjgillot and unassigned scottmcm Jun 17, 2023

rustbot added S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. and removed S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. labels Jun 17, 2023

rustbot added S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. and removed S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jun 19, 2023

cjgillot reviewed Jun 20, 2023

View reviewed changes

rustbot added S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. and removed S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. labels Jun 20, 2023

rustbot added S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. and removed S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jun 20, 2023

cjgillot reviewed Jun 23, 2023

View reviewed changes

rustbot added S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. and removed S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. labels Jun 24, 2023

rustbot removed the S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. label Jun 27, 2023

rustbot added the S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. label Jun 27, 2023

This comment has been minimized.

Sign in to view

Implement a SimplifyStaticSwitch MIR pass which

8f4ada4

removes unnecessary switches on statically known discriminants.

cjgillot reviewed Jul 4, 2023

View reviewed changes

cjgillot added S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. and removed S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. labels Jul 4, 2023

JohnCSimon closed this Dec 17, 2023

rustbot added the S-inactive Status: Inactive and waiting on the author. This is often applied to closed PRs. label Dec 17, 2023

	debug!(body = %tcx.def_path_str(body.source.def_id()));
	debug!(body = ?body.source.def_id());

	let predecessors: &IndexSlice<_, _> = body.basic_blocks.predecessors();
	let predecessors = body.basic_blocks.predecessors();

	let predecessors: &[BasicBlock] = &predecessors[block];
	let predecessors = &predecessors[block];

		"{pred:?}: {place:?} = {}::{variant:?}; goto -> {block:?};",
		tcx.def_path_str(def_id)

	"{pred:?}: {place:?} = {}::{variant:?}; goto -> {block:?};",
	tcx.def_path_str(def_id)
	"{pred:?}: {place:?} = {def_id:?}::{variant:?}; goto -> {block:?};"

	let old_goto = &mut data.terminator.as_mut().unwrap().kind;
	let old_goto = &mut data.terminator_mut().kind;

Simplify switches on a statically known discriminant in MIR. #112688

Simplify switches on a statically known discriminant in MIR. #112688

Conversation

JohnBobbo96 commented Jun 16, 2023 • edited Loading

rustbot commented Jun 16, 2023

bors commented Jun 16, 2023

scottmcm commented Jun 16, 2023

cjgillot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JohnBobbo96 commented Jun 17, 2023

JohnBobbo96 commented Jun 19, 2023 • edited Loading

JohnBobbo96 commented Jun 19, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JohnBobbo96 commented Jun 20, 2023

JohnBobbo96 commented Jun 20, 2023

cjgillot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JohnBobbo96 Jun 27, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JohnBobbo96 commented Jun 24, 2023

JohnBobbo96 commented Jun 27, 2023 • edited Loading

JohnBobbo96 commented Jun 27, 2023

This comment has been minimized.

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bors commented Jul 14, 2023

JohnCSimon commented Aug 27, 2023

JohnCSimon commented Dec 17, 2023

JohnBobbo96 commented Jun 16, 2023 •

edited

Loading

JohnBobbo96 commented Jun 19, 2023 •

edited

Loading

JohnBobbo96 Jun 27, 2023 •

edited

Loading

JohnBobbo96 commented Jun 27, 2023 •

edited

Loading