Auto merge of rust-lang#126578 - scottmcm:inlining-bonuses-too, r=davidtwco

bors · bors · commit 7a08f84627ff · 2024-06-21T02:06:27.000Z
Account for things that optimize out in inlining costs

This updates the MIR inlining `CostChecker` to have both bonuses and penalties, rather than just penalties.

That lets us add bonuses for some things where we want to encourage inlining without risking wrapping into a gigantic cost.  For example, `switchInt(const …)` we give an inlining bonus because codegen will actually eliminate the branch (and associated dead blocks) once it's monomorphized, so measuring both sides of the branch gives an unrealistically-high cost to it.  Similarly, an `unreachable` terminator gets a small bonus, because whatever branch leads there doesn't actually exist post-codegen.
diff --git a/compiler/rustc_mir_transform/src/cost_checker.rs b/compiler/rustc_mir_transform/src/cost_checker.rs
@@ -1,3 +1,4 @@
+use rustc_middle::bug;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, ParamEnv, Ty, TyCtxt};
@@ -6,13 +7,16 @@ const INSTR_COST: usize = 5;
 const CALL_PENALTY: usize = 25;
 const LANDINGPAD_PENALTY: usize = 50;
 const RESUME_PENALTY: usize = 45;
+const LARGE_SWITCH_PENALTY: usize = 20;
+const CONST_SWITCH_BONUS: usize = 10;
 
 /// Verify that the callee body is compatible with the caller.
 #[derive(Clone)]
 pub(crate) struct CostChecker<'b, 'tcx> {
     tcx: TyCtxt<'tcx>,
     param_env: ParamEnv<'tcx>,
-    cost: usize,
+    penalty: usize,
+    bonus: usize,
     callee_body: &'b Body<'tcx>,
     instance: Option<ty::Instance<'tcx>>,
 }
@@ -24,11 +28,11 @@ impl<'b, 'tcx> CostChecker<'b, 'tcx> {
         instance: Option<ty::Instance<'tcx>>,
         callee_body: &'b Body<'tcx>,
     ) -> CostChecker<'b, 'tcx> {
-        CostChecker { tcx, param_env, callee_body, instance, cost: 0 }
+        CostChecker { tcx, param_env, callee_body, instance, penalty: 0, bonus: 0 }
     }
 
     pub fn cost(&self) -> usize {
-        self.cost
+        usize::saturating_sub(self.penalty, self.bonus)
     }
 
     fn instantiate_ty(&self, v: Ty<'tcx>) -> Ty<'tcx> {
@@ -41,60 +45,100 @@ impl<'b, 'tcx> CostChecker<'b, 'tcx> {
 }
 
 impl<'tcx> Visitor<'tcx> for CostChecker<'_, 'tcx> {
-    fn visit_statement(&mut self, statement: &Statement<'tcx>, _: Location) {
-        // Don't count StorageLive/StorageDead in the inlining cost.
+    fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
+        // Most costs are in rvalues and terminators, not in statements.
         match statement.kind {
-            StatementKind::StorageLive(_)
-            | StatementKind::StorageDead(_)
-            | StatementKind::Deinit(_)
-            | StatementKind::Nop => {}
-            _ => self.cost += INSTR_COST,
+            StatementKind::Intrinsic(ref ndi) => {
+                self.penalty += match **ndi {
+                    NonDivergingIntrinsic::Assume(..) => INSTR_COST,
+                    NonDivergingIntrinsic::CopyNonOverlapping(..) => CALL_PENALTY,
+                };
+            }
+            _ => self.super_statement(statement, location),
+        }
+    }
+
+    fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, _location: Location) {
+        match rvalue {
+            Rvalue::NullaryOp(NullOp::UbChecks, ..) if !self.tcx.sess.ub_checks() => {
+                // If this is in optimized MIR it's because it's used later,
+                // so if we don't need UB checks this session, give a bonus
+                // here to offset the cost of the call later.
+                self.bonus += CALL_PENALTY;
+            }
+            // These are essentially constants that didn't end up in an Operand,
+            // so treat them as also being free.
+            Rvalue::NullaryOp(..) => {}
+            _ => self.penalty += INSTR_COST,
         }
     }
 
     fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, _: Location) {
-        let tcx = self.tcx;
-        match terminator.kind {
-            TerminatorKind::Drop { ref place, unwind, .. } => {
+        match &terminator.kind {
+            TerminatorKind::Drop { place, unwind, .. } => {
                 // If the place doesn't actually need dropping, treat it like a regular goto.
-                let ty = self.instantiate_ty(place.ty(self.callee_body, tcx).ty);
-                if ty.needs_drop(tcx, self.param_env) {
-                    self.cost += CALL_PENALTY;
+                let ty = self.instantiate_ty(place.ty(self.callee_body, self.tcx).ty);
+                if ty.needs_drop(self.tcx, self.param_env) {
+                    self.penalty += CALL_PENALTY;
                     if let UnwindAction::Cleanup(_) = unwind {
-                        self.cost += LANDINGPAD_PENALTY;
+                        self.penalty += LANDINGPAD_PENALTY;
                     }
-                } else {
-                    self.cost += INSTR_COST;
                 }
             }
-            TerminatorKind::Call { func: Operand::Constant(ref f), unwind, .. } => {
-                let fn_ty = self.instantiate_ty(f.const_.ty());
-                self.cost += if let ty::FnDef(def_id, _) = *fn_ty.kind()
-                    && tcx.intrinsic(def_id).is_some()
+            TerminatorKind::Call { func, unwind, .. } => {
+                self.penalty += if let Some((def_id, ..)) = func.const_fn_def()
+                    && self.tcx.intrinsic(def_id).is_some()
                 {
                     // Don't give intrinsics the extra penalty for calls
                     INSTR_COST
                 } else {
                     CALL_PENALTY
                 };
                 if let UnwindAction::Cleanup(_) = unwind {
-                    self.cost += LANDINGPAD_PENALTY;
+                    self.penalty += LANDINGPAD_PENALTY;
+                }
+            }
+            TerminatorKind::SwitchInt { discr, targets } => {
+                if discr.constant().is_some() {
+                    // Not only will this become a `Goto`, but likely other
+                    // things will be removable as unreachable.
+                    self.bonus += CONST_SWITCH_BONUS;
+                } else if targets.all_targets().len() > 3 {
+                    // More than false/true/unreachable gets extra cost.
+                    self.penalty += LARGE_SWITCH_PENALTY;
+                } else {
+                    self.penalty += INSTR_COST;
                 }
             }
-            TerminatorKind::Assert { unwind, .. } => {
-                self.cost += CALL_PENALTY;
+            TerminatorKind::Assert { unwind, msg, .. } => {
+                self.penalty +=
+                    if msg.is_optional_overflow_check() && !self.tcx.sess.overflow_checks() {
+                        INSTR_COST
+                    } else {
+                        CALL_PENALTY
+                    };
                 if let UnwindAction::Cleanup(_) = unwind {
-                    self.cost += LANDINGPAD_PENALTY;
+                    self.penalty += LANDINGPAD_PENALTY;
                 }
             }
-            TerminatorKind::UnwindResume => self.cost += RESUME_PENALTY,
+            TerminatorKind::UnwindResume => self.penalty += RESUME_PENALTY,
             TerminatorKind::InlineAsm { unwind, .. } => {
-                self.cost += INSTR_COST;
+                self.penalty += INSTR_COST;
                 if let UnwindAction::Cleanup(_) = unwind {
-                    self.cost += LANDINGPAD_PENALTY;
+                    self.penalty += LANDINGPAD_PENALTY;
                 }
             }
-            _ => self.cost += INSTR_COST,
+            TerminatorKind::Unreachable => {
+                self.bonus += INSTR_COST;
+            }
+            TerminatorKind::Goto { .. } | TerminatorKind::Return => {}
+            TerminatorKind::UnwindTerminate(..) => {}
+            kind @ (TerminatorKind::FalseUnwind { .. }
+            | TerminatorKind::FalseEdge { .. }
+            | TerminatorKind::Yield { .. }
+            | TerminatorKind::CoroutineDrop) => {
+                bug!("{kind:?} should not be in runtime MIR");
+            }
         }
     }
 }
diff --git a/library/core/src/slice/iter/macros.rs b/library/core/src/slice/iter/macros.rs
@@ -103,7 +103,8 @@ macro_rules! iterator {
                 // so this new pointer is inside `self` and thus guaranteed to be non-null.
                 unsafe {
                     if_zst!(mut self,
-                        len => *len = len.unchecked_sub(offset),
+                        // Using the intrinsic directly avoids emitting a UbCheck
+                        len => *len = crate::intrinsics::unchecked_sub(*len, offset),
                         _end => self.ptr = self.ptr.add(offset),
                     );
                 }
@@ -119,7 +120,8 @@ macro_rules! iterator {
                     // SAFETY: By our precondition, `offset` can be at most the
                     // current length, so the subtraction can never overflow.
                     len => unsafe {
-                        *len = len.unchecked_sub(offset);
+                        // Using the intrinsic directly avoids emitting a UbCheck
+                        *len = crate::intrinsics::unchecked_sub(*len, offset);
                         self.ptr
                     },
                     // SAFETY: the caller guarantees that `offset` doesn't exceed `self.len()`,
diff --git a/tests/codegen/issues/issue-112509-slice-get-andthen-get.rs b/tests/codegen/issues/issue-112509-slice-get-andthen-get.rs
@@ -3,8 +3,12 @@
 
 // CHECK-LABEL: @write_u8_variant_a
 // CHECK-NEXT: {{.*}}:
-// CHECK-NEXT: getelementptr
 // CHECK-NEXT: icmp ugt
+// CHECK-NEXT: getelementptr
+// CHECK-NEXT: select i1 {{.+}} null
+// CHECK-NEXT: insertvalue
+// CHECK-NEXT: insertvalue
+// CHECK-NEXT: ret
 #[no_mangle]
 pub fn write_u8_variant_a(bytes: &mut [u8], buf: u8, offset: usize) -> Option<&mut [u8]> {
     let buf = buf.to_le_bytes();
diff --git a/tests/crashes/123893.rs b/tests/crashes/123893.rs
@@ -11,10 +11,14 @@ fn generic_impl<T>() {
     impl<T> MagicTrait for T {
         const IS_BIG: bool = true;
     }
+    more_cost();
     if T::IS_BIG {
         big_impl::<i32>();
     }
 }
 
 #[inline(never)]
 fn big_impl<T>() {}
+
+#[inline(never)]
+fn more_cost() {}
diff --git a/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-abort.mir b/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-abort.mir
@@ -4,12 +4,87 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
     debug x => _1;
     debug n => _2;
     let mut _0: u16;
+    scope 1 (inlined <u16 as Step>::forward) {
+        let mut _8: u16;
+        scope 2 {
+        }
+        scope 3 (inlined <u16 as Step>::forward_checked) {
+            scope 4 {
+                scope 6 (inlined core::num::<impl u16>::checked_add) {
+                    let mut _7: bool;
+                    scope 7 {
+                    }
+                    scope 8 (inlined core::num::<impl u16>::overflowing_add) {
+                        let mut _5: (u16, bool);
+                        let _6: bool;
+                        scope 9 {
+                        }
+                    }
+                }
+            }
+            scope 5 (inlined convert::num::ptr_try_from_impls::<impl TryFrom<usize> for u16>::try_from) {
+                let mut _3: bool;
+                let mut _4: u16;
+            }
+        }
+        scope 10 (inlined Option::<u16>::is_none) {
+            scope 11 (inlined Option::<u16>::is_some) {
+            }
+        }
+        scope 12 (inlined core::num::<impl u16>::wrapping_add) {
+        }
+    }
 
     bb0: {
-        _0 = <u16 as Step>::forward(move _1, move _2) -> [return: bb1, unwind unreachable];
+        StorageLive(_4);
+        StorageLive(_3);
+        _3 = Gt(_2, const 65535_usize);
+        switchInt(move _3) -> [0: bb1, otherwise: bb5];
     }
 
     bb1: {
+        _4 = _2 as u16 (IntToInt);
+        StorageDead(_3);
+        StorageLive(_6);
+        StorageLive(_5);
+        _5 = AddWithOverflow(_1, _4);
+        _6 = (_5.1: bool);
+        StorageDead(_5);
+        StorageLive(_7);
+        _7 = unlikely(move _6) -> [return: bb2, unwind unreachable];
+    }
+
+    bb2: {
+        switchInt(move _7) -> [0: bb3, otherwise: bb4];
+    }
+
+    bb3: {
+        StorageDead(_7);
+        StorageDead(_6);
+        goto -> bb7;
+    }
+
+    bb4: {
+        StorageDead(_7);
+        StorageDead(_6);
+        goto -> bb6;
+    }
+
+    bb5: {
+        StorageDead(_3);
+        goto -> bb6;
+    }
+
+    bb6: {
+        assert(!const true, "attempt to compute `{} + {}`, which would overflow", const core::num::<impl u16>::MAX, const 1_u16) -> [success: bb7, unwind unreachable];
+    }
+
+    bb7: {
+        StorageLive(_8);
+        _8 = _2 as u16 (IntToInt);
+        _0 = Add(_1, _8);
+        StorageDead(_8);
+        StorageDead(_4);
         return;
     }
 }
diff --git a/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir b/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir
@@ -4,12 +4,87 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
     debug x => _1;
     debug n => _2;
     let mut _0: u16;
+    scope 1 (inlined <u16 as Step>::forward) {
+        let mut _8: u16;
+        scope 2 {
+        }
+        scope 3 (inlined <u16 as Step>::forward_checked) {
+            scope 4 {
+                scope 6 (inlined core::num::<impl u16>::checked_add) {
+                    let mut _7: bool;
+                    scope 7 {
+                    }
+                    scope 8 (inlined core::num::<impl u16>::overflowing_add) {
+                        let mut _5: (u16, bool);
+                        let _6: bool;
+                        scope 9 {
+                        }
+                    }
+                }
+            }
+            scope 5 (inlined convert::num::ptr_try_from_impls::<impl TryFrom<usize> for u16>::try_from) {
+                let mut _3: bool;
+                let mut _4: u16;
+            }
+        }
+        scope 10 (inlined Option::<u16>::is_none) {
+            scope 11 (inlined Option::<u16>::is_some) {
+            }
+        }
+        scope 12 (inlined core::num::<impl u16>::wrapping_add) {
+        }
+    }
 
     bb0: {
-        _0 = <u16 as Step>::forward(move _1, move _2) -> [return: bb1, unwind continue];
+        StorageLive(_4);
+        StorageLive(_3);
+        _3 = Gt(_2, const 65535_usize);
+        switchInt(move _3) -> [0: bb1, otherwise: bb5];
     }
 
     bb1: {
+        _4 = _2 as u16 (IntToInt);
+        StorageDead(_3);
+        StorageLive(_6);
+        StorageLive(_5);
+        _5 = AddWithOverflow(_1, _4);
+        _6 = (_5.1: bool);
+        StorageDead(_5);
+        StorageLive(_7);
+        _7 = unlikely(move _6) -> [return: bb2, unwind unreachable];
+    }
+
+    bb2: {
+        switchInt(move _7) -> [0: bb3, otherwise: bb4];
+    }
+
+    bb3: {
+        StorageDead(_7);
+        StorageDead(_6);
+        goto -> bb7;
+    }
+
+    bb4: {
+        StorageDead(_7);
+        StorageDead(_6);
+        goto -> bb6;
+    }
+
+    bb5: {
+        StorageDead(_3);
+        goto -> bb6;
+    }
+
+    bb6: {
+        assert(!const true, "attempt to compute `{} + {}`, which would overflow", const core::num::<impl u16>::MAX, const 1_u16) -> [success: bb7, unwind continue];
+    }
+
+    bb7: {
+        StorageLive(_8);
+        _8 = _2 as u16 (IntToInt);
+        _0 = Add(_1, _8);
+        StorageDead(_8);
+        StorageDead(_4);
         return;
     }
 }
diff --git a/tests/mir-opt/pre-codegen/loops.mapped.PreCodegen.after.mir b/tests/mir-opt/pre-codegen/loops.mapped.PreCodegen.after.mir
diff --git a/tests/mir-opt/pre-codegen/slice_iter.enumerated_loop.PreCodegen.after.panic-abort.mir b/tests/mir-opt/pre-codegen/slice_iter.enumerated_loop.PreCodegen.after.panic-abort.mir

Original file line number	Diff line number	Diff line change
`@@ -11,10 +11,14 @@ fn generic_impl<T>() {`
`11`	`11`	`impl<T> MagicTrait for T {`
`12`	`12`	`const IS_BIG: bool = true;`
`13`	`13`	`}`
	`14`	`+ more_cost();`
`14`	`15`	`if T::IS_BIG {`
`15`	`16`	`big_impl::<i32>();`
`16`	`17`	`}`
`17`	`18`	`}`
`18`	`19`
`19`	`20`	`#[inline(never)]`
`20`	`21`	`fn big_impl<T>() {}`
	`22`	`+`
	`23`	`+#[inline(never)]`
	`24`	`+fn more_cost() {}`