From e3a9eaf928a2028ac383e3025fe4e2f4ca919fda Mon Sep 17 00:00:00 2001 From: DianQK Date: Wed, 18 Sep 2024 21:41:37 +0800 Subject: [PATCH] Apply `EarlyOtherwiseBranch` to scalar value --- .../src/early_otherwise_branch.rs | 260 ++++++++++++------ ...wise_branch.opt5.EarlyOtherwiseBranch.diff | 77 ++++++ ...anch.opt5_failed.EarlyOtherwiseBranch.diff | 61 ++++ ...opt5_failed_type.EarlyOtherwiseBranch.diff | 61 ++++ tests/mir-opt/early_otherwise_branch.rs | 48 ++++ 5 files changed, 422 insertions(+), 85 deletions(-) create mode 100644 tests/mir-opt/early_otherwise_branch.opt5.EarlyOtherwiseBranch.diff create mode 100644 tests/mir-opt/early_otherwise_branch.opt5_failed.EarlyOtherwiseBranch.diff create mode 100644 tests/mir-opt/early_otherwise_branch.opt5_failed_type.EarlyOtherwiseBranch.diff diff --git a/compiler/rustc_mir_transform/src/early_otherwise_branch.rs b/compiler/rustc_mir_transform/src/early_otherwise_branch.rs index 704ed508b22..aed7f20aaea 100644 --- a/compiler/rustc_mir_transform/src/early_otherwise_branch.rs +++ b/compiler/rustc_mir_transform/src/early_otherwise_branch.rs @@ -133,18 +133,29 @@ impl<'tcx> crate::MirPass<'tcx> for EarlyOtherwiseBranch { let mut patch = MirPatch::new(body); - // create temp to store second discriminant in, `_s` in example above - let second_discriminant_temp = - patch.new_temp(opt_data.child_ty, opt_data.child_source.span); + let (second_discriminant_temp, second_operand) = if opt_data.need_hoist_discriminant { + // create temp to store second discriminant in, `_s` in example above + let second_discriminant_temp = + patch.new_temp(opt_data.child_ty, opt_data.child_source.span); - patch.add_statement(parent_end, StatementKind::StorageLive(second_discriminant_temp)); + patch.add_statement( + parent_end, + StatementKind::StorageLive(second_discriminant_temp), + ); - // create assignment of discriminant - patch.add_assign( - parent_end, - Place::from(second_discriminant_temp), - Rvalue::Discriminant(opt_data.child_place), - ); + // create assignment of discriminant + patch.add_assign( + parent_end, + Place::from(second_discriminant_temp), + Rvalue::Discriminant(opt_data.child_place), + ); + ( + Some(second_discriminant_temp), + Operand::Move(Place::from(second_discriminant_temp)), + ) + } else { + (None, Operand::Copy(opt_data.child_place)) + }; // create temp to store inequality comparison between the two discriminants, `_t` in // example above @@ -153,11 +164,9 @@ impl<'tcx> crate::MirPass<'tcx> for EarlyOtherwiseBranch { let comp_temp = patch.new_temp(comp_res_type, opt_data.child_source.span); patch.add_statement(parent_end, StatementKind::StorageLive(comp_temp)); - // create inequality comparison between the two discriminants - let comp_rvalue = Rvalue::BinaryOp( - nequal, - Box::new((parent_op.clone(), Operand::Move(Place::from(second_discriminant_temp)))), - ); + // create inequality comparison + let comp_rvalue = + Rvalue::BinaryOp(nequal, Box::new((parent_op.clone(), second_operand))); patch.add_statement( parent_end, StatementKind::Assign(Box::new((Place::from(comp_temp), comp_rvalue))), @@ -193,8 +202,13 @@ impl<'tcx> crate::MirPass<'tcx> for EarlyOtherwiseBranch { TerminatorKind::if_(Operand::Move(Place::from(comp_temp)), true_case, false_case), ); - // generate StorageDead for the second_discriminant_temp not in use anymore - patch.add_statement(parent_end, StatementKind::StorageDead(second_discriminant_temp)); + if let Some(second_discriminant_temp) = second_discriminant_temp { + // generate StorageDead for the second_discriminant_temp not in use anymore + patch.add_statement( + parent_end, + StatementKind::StorageDead(second_discriminant_temp), + ); + } // Generate a StorageDead for comp_temp in each of the targets, since we moved it into // the switch @@ -222,6 +236,7 @@ struct OptimizationData<'tcx> { child_place: Place<'tcx>, child_ty: Ty<'tcx>, child_source: SourceInfo, + need_hoist_discriminant: bool, } fn evaluate_candidate<'tcx>( @@ -235,44 +250,12 @@ fn evaluate_candidate<'tcx>( return None; }; let parent_ty = parent_discr.ty(body.local_decls(), tcx); - if !bbs[targets.otherwise()].is_empty_unreachable() { - // Someone could write code like this: - // ```rust - // let Q = val; - // if discriminant(P) == otherwise { - // let ptr = &mut Q as *mut _ as *mut u8; - // // It may be difficult for us to effectively determine whether values are valid. - // // Invalid values can come from all sorts of corners. - // unsafe { *ptr = 10; } - // } - // - // match P { - // A => match Q { - // A => { - // // code - // } - // _ => { - // // don't use Q - // } - // } - // _ => { - // // don't use Q - // } - // }; - // ``` - // - // Hoisting the `discriminant(Q)` out of the `A` arm causes us to compute the discriminant - // of an invalid value, which is UB. - // In order to fix this, **we would either need to show that the discriminant computation of - // `place` is computed in all branches**. - // FIXME(#95162) For the moment, we adopt a conservative approach and - // consider only the `otherwise` branch has no statements and an unreachable terminator. - return None; - } let (_, child) = targets.iter().next()?; - let child_terminator = &bbs[child].terminator(); - let TerminatorKind::SwitchInt { targets: child_targets, discr: child_discr } = - &child_terminator.kind + + let Terminator { + kind: TerminatorKind::SwitchInt { targets: child_targets, discr: child_discr }, + source_info, + } = bbs[child].terminator() else { return None; }; @@ -280,25 +263,115 @@ fn evaluate_candidate<'tcx>( if child_ty != parent_ty { return None; } - let Some(StatementKind::Assign(boxed)) = &bbs[child].statements.first().map(|x| &x.kind) else { + + // We only handle: + // ``` + // bb4: { + // _8 = discriminant((_3.1: Enum1)); + // switchInt(move _8) -> [2: bb7, otherwise: bb1]; + // } + // ``` + // and + // ``` + // bb2: { + // switchInt((_3.1: u64)) -> [1: bb5, otherwise: bb1]; + // } + // ``` + if bbs[child].statements.len() > 1 { return None; + } + + // When thie BB has exactly one statement, this statement should be discriminant. + let need_hoist_discriminant = bbs[child].statements.len() == 1; + let child_place = if need_hoist_discriminant { + if !bbs[targets.otherwise()].is_empty_unreachable() { + // Someone could write code like this: + // ```rust + // let Q = val; + // if discriminant(P) == otherwise { + // let ptr = &mut Q as *mut _ as *mut u8; + // // It may be difficult for us to effectively determine whether values are valid. + // // Invalid values can come from all sorts of corners. + // unsafe { *ptr = 10; } + // } + // + // match P { + // A => match Q { + // A => { + // // code + // } + // _ => { + // // don't use Q + // } + // } + // _ => { + // // don't use Q + // } + // }; + // ``` + // + // Hoisting the `discriminant(Q)` out of the `A` arm causes us to compute the discriminant of an + // invalid value, which is UB. + // In order to fix this, **we would either need to show that the discriminant computation of + // `place` is computed in all branches**. + // FIXME(#95162) For the moment, we adopt a conservative approach and + // consider only the `otherwise` branch has no statements and an unreachable terminator. + return None; + } + // Handle: + // ``` + // bb4: { + // _8 = discriminant((_3.1: Enum1)); + // switchInt(move _8) -> [2: bb7, otherwise: bb1]; + // } + // ``` + let [ + Statement { + kind: StatementKind::Assign(box (_, Rvalue::Discriminant(child_place))), + .. + }, + ] = bbs[child].statements.as_slice() + else { + return None; + }; + *child_place + } else { + // Handle: + // ``` + // bb2: { + // switchInt((_3.1: u64)) -> [1: bb5, otherwise: bb1]; + // } + // ``` + let Operand::Copy(child_place) = child_discr else { + return None; + }; + *child_place }; - let (_, Rvalue::Discriminant(child_place)) = &**boxed else { - return None; + let destination = if need_hoist_discriminant || bbs[targets.otherwise()].is_empty_unreachable() + { + child_targets.otherwise() + } else { + targets.otherwise() }; - let destination = child_targets.otherwise(); // Verify that the optimization is legal for each branch for (value, child) in targets.iter() { - if !verify_candidate_branch(&bbs[child], value, *child_place, destination) { + if !verify_candidate_branch( + &bbs[child], + value, + child_place, + destination, + need_hoist_discriminant, + ) { return None; } } Some(OptimizationData { destination, - child_place: *child_place, + child_place, child_ty, - child_source: child_terminator.source_info, + child_source: *source_info, + need_hoist_discriminant, }) } @@ -307,31 +380,48 @@ fn verify_candidate_branch<'tcx>( value: u128, place: Place<'tcx>, destination: BasicBlock, + need_hoist_discriminant: bool, ) -> bool { - // In order for the optimization to be correct, the branch must... - // ...have exactly one statement - if let [statement] = branch.statements.as_slice() - // ...assign the discriminant of `place` in that statement - && let StatementKind::Assign(boxed) = &statement.kind - && let (discr_place, Rvalue::Discriminant(from_place)) = &**boxed - && *from_place == place - // ...make that assignment to a local - && discr_place.projection.is_empty() - // ...terminate on a `SwitchInt` that invalidates that local - && let TerminatorKind::SwitchInt { discr: switch_op, targets, .. } = - &branch.terminator().kind - && *switch_op == Operand::Move(*discr_place) - // ...fall through to `destination` if the switch misses - && destination == targets.otherwise() - // ...have a branch for value `value` - && let mut iter = targets.iter() - && let Some((target_value, _)) = iter.next() - && target_value == value - // ...and have no more branches - && iter.next().is_none() - { - true + // In order for the optimization to be correct, the terminator must be a `SwitchInt`. + let TerminatorKind::SwitchInt { discr: switch_op, targets } = &branch.terminator().kind else { + return false; + }; + if need_hoist_discriminant { + // If we need hoist discriminant, the branch must have exactly one statement. + let [statement] = branch.statements.as_slice() else { + return false; + }; + // The statement must assign the discriminant of `place`. + let StatementKind::Assign(box (discr_place, Rvalue::Discriminant(from_place))) = + statement.kind + else { + return false; + }; + if from_place != place { + return false; + } + // The assignment must invalidate a local that terminate on a `SwitchInt`. + if !discr_place.projection.is_empty() || *switch_op != Operand::Move(discr_place) { + return false; + } } else { - false + // If we don't need hoist discriminant, the branch must not have any statements. + if !branch.statements.is_empty() { + return false; + } + // The place on `SwitchInt` must be the same. + if *switch_op != Operand::Copy(place) { + return false; + } } + // It must fall through to `destination` if the switch misses. + if destination != targets.otherwise() { + return false; + } + // It must have exactly one branch for value `value` and have no more branches. + let mut iter = targets.iter(); + let (Some((target_value, _)), None) = (iter.next(), iter.next()) else { + return false; + }; + target_value == value } diff --git a/tests/mir-opt/early_otherwise_branch.opt5.EarlyOtherwiseBranch.diff b/tests/mir-opt/early_otherwise_branch.opt5.EarlyOtherwiseBranch.diff new file mode 100644 index 00000000000..b7447ef0c46 --- /dev/null +++ b/tests/mir-opt/early_otherwise_branch.opt5.EarlyOtherwiseBranch.diff @@ -0,0 +1,77 @@ +- // MIR for `opt5` before EarlyOtherwiseBranch ++ // MIR for `opt5` after EarlyOtherwiseBranch + + fn opt5(_1: u32, _2: u32) -> u32 { + debug x => _1; + debug y => _2; + let mut _0: u32; + let mut _3: (u32, u32); + let mut _4: u32; + let mut _5: u32; ++ let mut _6: bool; + + bb0: { + StorageLive(_3); + StorageLive(_4); + _4 = copy _1; + StorageLive(_5); + _5 = copy _2; + _3 = (move _4, move _5); + StorageDead(_5); + StorageDead(_4); +- switchInt(copy (_3.0: u32)) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1]; ++ StorageLive(_6); ++ _6 = Ne(copy (_3.0: u32), copy (_3.1: u32)); ++ switchInt(move _6) -> [0: bb6, otherwise: bb1]; + } + + bb1: { ++ StorageDead(_6); + _0 = const 0_u32; +- goto -> bb8; ++ goto -> bb5; + } + + bb2: { +- switchInt(copy (_3.1: u32)) -> [1: bb7, otherwise: bb1]; ++ _0 = const 6_u32; ++ goto -> bb5; + } + + bb3: { +- switchInt(copy (_3.1: u32)) -> [2: bb6, otherwise: bb1]; ++ _0 = const 5_u32; ++ goto -> bb5; + } + + bb4: { +- switchInt(copy (_3.1: u32)) -> [3: bb5, otherwise: bb1]; ++ _0 = const 4_u32; ++ goto -> bb5; + } + + bb5: { +- _0 = const 6_u32; +- goto -> bb8; ++ StorageDead(_3); ++ return; + } + + bb6: { +- _0 = const 5_u32; +- goto -> bb8; +- } +- +- bb7: { +- _0 = const 4_u32; +- goto -> bb8; +- } +- +- bb8: { +- StorageDead(_3); +- return; ++ StorageDead(_6); ++ switchInt(copy (_3.0: u32)) -> [1: bb4, 2: bb3, 3: bb2, otherwise: bb1]; + } + } + diff --git a/tests/mir-opt/early_otherwise_branch.opt5_failed.EarlyOtherwiseBranch.diff b/tests/mir-opt/early_otherwise_branch.opt5_failed.EarlyOtherwiseBranch.diff new file mode 100644 index 00000000000..af16271f8b1 --- /dev/null +++ b/tests/mir-opt/early_otherwise_branch.opt5_failed.EarlyOtherwiseBranch.diff @@ -0,0 +1,61 @@ +- // MIR for `opt5_failed` before EarlyOtherwiseBranch ++ // MIR for `opt5_failed` after EarlyOtherwiseBranch + + fn opt5_failed(_1: u32, _2: u32) -> u32 { + debug x => _1; + debug y => _2; + let mut _0: u32; + let mut _3: (u32, u32); + let mut _4: u32; + let mut _5: u32; + + bb0: { + StorageLive(_3); + StorageLive(_4); + _4 = copy _1; + StorageLive(_5); + _5 = copy _2; + _3 = (move _4, move _5); + StorageDead(_5); + StorageDead(_4); + switchInt(copy (_3.0: u32)) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1]; + } + + bb1: { + _0 = const 0_u32; + goto -> bb8; + } + + bb2: { + switchInt(copy (_3.1: u32)) -> [1: bb7, otherwise: bb1]; + } + + bb3: { + switchInt(copy (_3.1: u32)) -> [2: bb6, otherwise: bb1]; + } + + bb4: { + switchInt(copy (_3.1: u32)) -> [2: bb5, otherwise: bb1]; + } + + bb5: { + _0 = const 6_u32; + goto -> bb8; + } + + bb6: { + _0 = const 5_u32; + goto -> bb8; + } + + bb7: { + _0 = const 4_u32; + goto -> bb8; + } + + bb8: { + StorageDead(_3); + return; + } + } + diff --git a/tests/mir-opt/early_otherwise_branch.opt5_failed_type.EarlyOtherwiseBranch.diff b/tests/mir-opt/early_otherwise_branch.opt5_failed_type.EarlyOtherwiseBranch.diff new file mode 100644 index 00000000000..23f14b843b3 --- /dev/null +++ b/tests/mir-opt/early_otherwise_branch.opt5_failed_type.EarlyOtherwiseBranch.diff @@ -0,0 +1,61 @@ +- // MIR for `opt5_failed_type` before EarlyOtherwiseBranch ++ // MIR for `opt5_failed_type` after EarlyOtherwiseBranch + + fn opt5_failed_type(_1: u32, _2: u64) -> u32 { + debug x => _1; + debug y => _2; + let mut _0: u32; + let mut _3: (u32, u64); + let mut _4: u32; + let mut _5: u64; + + bb0: { + StorageLive(_3); + StorageLive(_4); + _4 = copy _1; + StorageLive(_5); + _5 = copy _2; + _3 = (move _4, move _5); + StorageDead(_5); + StorageDead(_4); + switchInt(copy (_3.0: u32)) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1]; + } + + bb1: { + _0 = const 0_u32; + goto -> bb8; + } + + bb2: { + switchInt(copy (_3.1: u64)) -> [1: bb7, otherwise: bb1]; + } + + bb3: { + switchInt(copy (_3.1: u64)) -> [2: bb6, otherwise: bb1]; + } + + bb4: { + switchInt(copy (_3.1: u64)) -> [3: bb5, otherwise: bb1]; + } + + bb5: { + _0 = const 6_u32; + goto -> bb8; + } + + bb6: { + _0 = const 5_u32; + goto -> bb8; + } + + bb7: { + _0 = const 4_u32; + goto -> bb8; + } + + bb8: { + StorageDead(_3); + return; + } + } + diff --git a/tests/mir-opt/early_otherwise_branch.rs b/tests/mir-opt/early_otherwise_branch.rs index 47bd4be295b..382c38ceb3a 100644 --- a/tests/mir-opt/early_otherwise_branch.rs +++ b/tests/mir-opt/early_otherwise_branch.rs @@ -78,9 +78,57 @@ fn opt4(x: Option2, y: Option2) -> u32 { } } +// EMIT_MIR early_otherwise_branch.opt5.EarlyOtherwiseBranch.diff +fn opt5(x: u32, y: u32) -> u32 { + // CHECK-LABEL: fn opt5( + // CHECK: let mut [[CMP_LOCAL:_.*]]: bool; + // CHECK: bb0: { + // CHECK: [[CMP_LOCAL]] = Ne( + // CHECK: switchInt(move [[CMP_LOCAL]]) -> [ + // CHECK-NEXT: } + match (x, y) { + (1, 1) => 4, + (2, 2) => 5, + (3, 3) => 6, + _ => 0, + } +} + +// EMIT_MIR early_otherwise_branch.opt5_failed.EarlyOtherwiseBranch.diff +fn opt5_failed(x: u32, y: u32) -> u32 { + // CHECK-LABEL: fn opt5_failed( + // CHECK: bb0: { + // CHECK-NOT: Ne( + // CHECK: switchInt( + // CHECK-NEXT: } + match (x, y) { + (1, 1) => 4, + (2, 2) => 5, + (3, 2) => 6, + _ => 0, + } +} + +// EMIT_MIR early_otherwise_branch.opt5_failed_type.EarlyOtherwiseBranch.diff +fn opt5_failed_type(x: u32, y: u64) -> u32 { + // CHECK-LABEL: fn opt5_failed_type( + // CHECK: bb0: { + // CHECK-NOT: Ne( + // CHECK: switchInt( + // CHECK-NEXT: } + match (x, y) { + (1, 1) => 4, + (2, 2) => 5, + (3, 3) => 6, + _ => 0, + } +} + fn main() { opt1(None, Some(0)); opt2(None, Some(0)); opt3(Option2::None, Option2::Some(false)); opt4(Option2::None, Option2::Some(0)); + opt5(0, 0); + opt5_failed(0, 0); }