Rollup merge of #79117 - cjkenn:mir-fuel, r=oli-obk

add optimization fuel checks to some mir passes

Fixes #77402

Inserts a bunch of calls to `consider_optimizing`. Note that `consider_optimizing` is the method that actually decrements the fuel count, so the point at which it's called is when the optimization takes place, from a fuel perspective. This means that where we call it has some thought behind it:

1. We probably don't want to decrement the fuel count before other simple checks, otherwise we count an optimization as being performed even if nothing was mutated (ie. it returned early).
2. In cases like `InstCombine`, where we gather optimizations in a pass and then mutate values, we probably would rather skip the gathering pass for performance reasons rather than skip the mutations afterwards.

Author: Dylan-DPC

compiler/rustc_mir/src/transform/const_prop.rs

@@ -800,7 +800,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
             }
         }
 
-        trace!("attepting to replace {:?} with {:?}", rval, value);
+        trace!("attempting to replace {:?} with {:?}", rval, value);
         if let Err(e) = self.ecx.const_validate_operand(
             value,
             vec![],
@@ -890,6 +890,10 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
             return false;
         }
 
+        if !self.tcx.consider_optimizing(|| format!("ConstantPropagation - OpTy: {:?}", op)) {
+            return false;
+        }
+
         match *op {
             interpret::Operand::Immediate(Immediate::Scalar(ScalarMaybeUninit::Scalar(s))) => {
                 s.is_bits()

compiler/rustc_mir/src/transform/early_otherwise_branch.rs

@@ -46,6 +46,10 @@ impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
         let should_cleanup = !opts_to_apply.is_empty();
 
         for opt_to_apply in opts_to_apply {
+            if !tcx.consider_optimizing(|| format!("EarlyOtherwiseBranch {:?}", &opt_to_apply)) {
+                break;
+            }
+
             trace!("SUCCESS: found optimization possibility to apply: {:?}", &opt_to_apply);
 
             let statements_before =

compiler/rustc_mir/src/transform/instcombine.rs

@@ -39,13 +39,21 @@ pub struct InstCombineVisitor<'tcx> {
     tcx: TyCtxt<'tcx>,
 }
 
+impl<'tcx> InstCombineVisitor<'tcx> {
+    fn should_combine(&self, rvalue: &Rvalue<'tcx>, location: Location) -> bool {
+        self.tcx.consider_optimizing(|| {
+            format!("InstCombine - Rvalue: {:?} Location: {:?}", rvalue, location)
+        })
+    }
+}
+
 impl<'tcx> MutVisitor<'tcx> for InstCombineVisitor<'tcx> {
     fn tcx(&self) -> TyCtxt<'tcx> {
         self.tcx
     }
 
     fn visit_rvalue(&mut self, rvalue: &mut Rvalue<'tcx>, location: Location) {
-        if self.optimizations.and_stars.remove(&location) {
+        if self.optimizations.and_stars.remove(&location) && self.should_combine(rvalue, location) {
             debug!("replacing `&*`: {:?}", rvalue);
             let new_place = match rvalue {
                 Rvalue::Ref(_, _, place) => {
@@ -67,18 +75,24 @@ impl<'tcx> MutVisitor<'tcx> for InstCombineVisitor<'tcx> {
         }
 
         if let Some(constant) = self.optimizations.arrays_lengths.remove(&location) {
-            debug!("replacing `Len([_; N])`: {:?}", rvalue);
-            *rvalue = Rvalue::Use(Operand::Constant(box constant));
+            if self.should_combine(rvalue, location) {
+                debug!("replacing `Len([_; N])`: {:?}", rvalue);
+                *rvalue = Rvalue::Use(Operand::Constant(box constant));
+            }
         }
 
         if let Some(operand) = self.optimizations.unneeded_equality_comparison.remove(&location) {
-            debug!("replacing {:?} with {:?}", rvalue, operand);
-            *rvalue = Rvalue::Use(operand);
+            if self.should_combine(rvalue, location) {
+                debug!("replacing {:?} with {:?}", rvalue, operand);
+                *rvalue = Rvalue::Use(operand);
+            }
         }
 
         if let Some(place) = self.optimizations.unneeded_deref.remove(&location) {
-            debug!("unneeded_deref: replacing {:?} with {:?}", rvalue, place);
-            *rvalue = Rvalue::Use(Operand::Copy(place));
+            if self.should_combine(rvalue, location) {
+                debug!("unneeded_deref: replacing {:?} with {:?}", rvalue, place);
+                *rvalue = Rvalue::Use(Operand::Copy(place));
+            }
         }
 
         self.super_rvalue(rvalue, location)

compiler/rustc_mir/src/transform/match_branches.rs

@@ -43,8 +43,13 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
         }
 
         let param_env = tcx.param_env(body.source.def_id());
+        let def_id = body.source.def_id();
         let (bbs, local_decls) = body.basic_blocks_and_local_decls_mut();
         'outer: for bb_idx in bbs.indices() {
+            if !tcx.consider_optimizing(|| format!("MatchBranchSimplification {:?} ", def_id)) {
+                continue;
+            }
+
             let (discr, val, switch_ty, first, second) = match bbs[bb_idx].terminator().kind {
                 TerminatorKind::SwitchInt {
                     discr: ref discr @ (Operand::Copy(_) | Operand::Move(_)),

compiler/rustc_mir/src/transform/multiple_return_terminators.rs

@@ -16,6 +16,7 @@ impl<'tcx> MirPass<'tcx> for MultipleReturnTerminators {
 
         // find basic blocks with no statement and a return terminator
         let mut bbs_simple_returns = BitSet::new_empty(body.basic_blocks().len());
+        let def_id = body.source.def_id();
         let bbs = body.basic_blocks_mut();
         for idx in bbs.indices() {
             if bbs[idx].statements.is_empty()
@@ -26,6 +27,10 @@ impl<'tcx> MirPass<'tcx> for MultipleReturnTerminators {
         }
 
         for bb in bbs {
+            if !tcx.consider_optimizing(|| format!("MultipleReturnTerminators {:?} ", def_id)) {
+                break;
+            }
+
             if let TerminatorKind::Goto { target } = bb.terminator().kind {
                 if bbs_simple_returns.contains(target) {
                     bb.terminator_mut().kind = TerminatorKind::Return;

compiler/rustc_mir/src/transform/nrvo.rs

@@ -38,18 +38,22 @@ impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
             return;
         }
 
+        let def_id = body.source.def_id();
         let returned_local = match local_eligible_for_nrvo(body) {
             Some(l) => l,
             None => {
-                debug!("`{:?}` was ineligible for NRVO", body.source.def_id());
+                debug!("`{:?}` was ineligible for NRVO", def_id);
                 return;
             }
         };
 
+        if !tcx.consider_optimizing(|| format!("RenameReturnPlace {:?}", def_id)) {
+            return;
+        }
+
         debug!(
             "`{:?}` was eligible for NRVO, making {:?} the return place",
-            body.source.def_id(),
-            returned_local
+            def_id, returned_local
         );
 
         RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body(body);

compiler/rustc_mir/src/transform/remove_unneeded_drops.rs

@@ -21,6 +21,12 @@ impl<'tcx> MirPass<'tcx> for RemoveUnneededDrops {
         opt_finder.visit_body(body);
         let should_simplify = !opt_finder.optimizations.is_empty();
         for (loc, target) in opt_finder.optimizations {
+            if !tcx
+                .consider_optimizing(|| format!("RemoveUnneededDrops {:?} ", body.source.def_id()))
+            {
+                break;
+            }
+
             let terminator = body.basic_blocks_mut()[loc.block].terminator_mut();
             debug!("SUCCESS: replacing `drop` with goto({:?})", target);
             terminator.kind = TerminatorKind::Goto { target };

compiler/rustc_mir/src/transform/unreachable_prop.rs

@@ -50,6 +50,12 @@ impl MirPass<'_> for UnreachablePropagation {
 
         let replaced = !replacements.is_empty();
         for (bb, terminator_kind) in replacements {
+            if !tcx.consider_optimizing(|| {
+                format!("UnreachablePropagation {:?} ", body.source.def_id())
+            }) {
+                break;
+            }
+
             body.basic_blocks_mut()[bb].terminator_mut().kind = terminator_kind;
         }
 

src/test/ui/print-fuel/print-fuel.rs

@@ -2,8 +2,8 @@
 #![allow(dead_code)]
 
 // (#55495: The --error-format is to sidestep an issue in our test harness)
-// compile-flags: --error-format human -Z print-fuel=foo
-// build-pass (FIXME(62277): could be check-pass?)
+// compile-flags: -C opt-level=0 --error-format human -Z print-fuel=foo
+// check-pass
 
 struct S1(u8, u16, u8);
 struct S2(u8, u16, u8);