Avoid MIR bloat in inlining

In 126578 we ended up with more binary size increases than expected.

This change attempts to avoid inlining large things into small things, to avoid that kind of increase, in cases when top-down inlining will still be able to do that inlining later.

Author: scottmcm

compiler/rustc_mir_transform/src/cost_checker.rs

@@ -31,6 +31,37 @@ impl<'b, 'tcx> CostChecker<'b, 'tcx> {
         CostChecker { tcx, param_env, callee_body, instance, penalty: 0, bonus: 0 }
     }
 
+    /// Add function-level costs not well-represented by the block-level costs.
+    ///
+    /// Needed because the `CostChecker` is used sometimes for just blocks,
+    /// and even the full `Inline` doesn't call `visit_body`, so there's nowhere
+    /// to put this logic in the visitor.
+    pub fn add_function_level_costs(&mut self) {
+        fn is_call_like(bbd: &BasicBlockData<'_>) -> bool {
+            use TerminatorKind::*;
+            match bbd.terminator().kind {
+                Call { .. } | Drop { .. } | Assert { .. } | InlineAsm { .. } => true,
+
+                Goto { .. }
+                | SwitchInt { .. }
+                | UnwindResume
+                | UnwindTerminate(_)
+                | Return
+                | Unreachable => false,
+
+                Yield { .. } | CoroutineDrop | FalseEdge { .. } | FalseUnwind { .. } => {
+                    unreachable!()
+                }
+            }
+        }
+
+        // If the only has one Call (or similar), inlining isn't increasing the total
+        // number of calls, so give extra encouragement to inlining that.
+        if self.callee_body.basic_blocks.iter().filter(|bbd| is_call_like(bbd)).count() == 1 {
+            self.bonus += CALL_PENALTY;
+        }
+    }
+
     pub fn cost(&self) -> usize {
         usize::saturating_sub(self.penalty, self.bonus)
     }

compiler/rustc_mir_transform/src/inline.rs

@@ -85,13 +85,18 @@ fn inline<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool {
     }
 
     let param_env = tcx.param_env_reveal_all_normalized(def_id);
+    let codegen_fn_attrs = tcx.codegen_fn_attrs(def_id);
 
     let mut this = Inliner {
         tcx,
         param_env,
-        codegen_fn_attrs: tcx.codegen_fn_attrs(def_id),
+        codegen_fn_attrs,
         history: Vec::new(),
         changed: false,
+        caller_is_inline_forwarder: matches!(
+            codegen_fn_attrs.inline,
+            InlineAttr::Hint | InlineAttr::Always
+        ) && body_is_forwarder(body),
     };
     let blocks = START_BLOCK..body.basic_blocks.next_index();
     this.process_blocks(body, blocks);
@@ -111,6 +116,9 @@ struct Inliner<'tcx> {
     history: Vec<DefId>,
     /// Indicates that the caller body has been modified.
     changed: bool,
+    /// Indicates that the caller is #[inline] and just calls another function,
+    /// and thus we can inline less into it as it'll be inlined itself.
+    caller_is_inline_forwarder: bool,
 }
 
 impl<'tcx> Inliner<'tcx> {
@@ -485,7 +493,9 @@ impl<'tcx> Inliner<'tcx> {
     ) -> Result<(), &'static str> {
         let tcx = self.tcx;
 
-        let mut threshold = if cross_crate_inlinable {
+        let mut threshold = if self.caller_is_inline_forwarder {
+            self.tcx.sess.opts.unstable_opts.inline_mir_forwarder_threshold.unwrap_or(30)
+        } else if cross_crate_inlinable {
             self.tcx.sess.opts.unstable_opts.inline_mir_hint_threshold.unwrap_or(100)
         } else {
             self.tcx.sess.opts.unstable_opts.inline_mir_threshold.unwrap_or(50)
@@ -504,6 +514,8 @@ impl<'tcx> Inliner<'tcx> {
         let mut checker =
             CostChecker::new(self.tcx, self.param_env, Some(callsite.callee), callee_body);
 
+        checker.add_function_level_costs();
+
         // Traverse the MIR manually so we can account for the effects of inlining on the CFG.
         let mut work_list = vec![START_BLOCK];
         let mut visited = BitSet::new_empty(callee_body.basic_blocks.len());
@@ -1091,3 +1103,37 @@ fn try_instance_mir<'tcx>(
     }
     Ok(tcx.instance_mir(instance))
 }
+
+fn body_is_forwarder(body: &Body<'_>) -> bool {
+    let TerminatorKind::Call { target, .. } = body.basic_blocks[START_BLOCK].terminator().kind
+    else {
+        return false;
+    };
+    if let Some(target) = target {
+        let TerminatorKind::Return = body.basic_blocks[target].terminator().kind else {
+            return false;
+        };
+    }
+
+    let max_blocks = if !body.is_polymorphic {
+        2
+    } else if target.is_none() {
+        3
+    } else {
+        4
+    };
+    if body.basic_blocks.len() > max_blocks {
+        return false;
+    }
+
+    body.basic_blocks.iter_enumerated().all(|(bb, bb_data)| {
+        bb == START_BLOCK
+            || matches!(
+                bb_data.terminator().kind,
+                TerminatorKind::Return
+                    | TerminatorKind::Drop { .. }
+                    | TerminatorKind::UnwindResume
+                    | TerminatorKind::UnwindTerminate(_)
+            )
+    })
+}

compiler/rustc_session/src/options.rs

@@ -1766,6 +1766,8 @@ options! {
         "enable LLVM inlining (default: yes)"),
     inline_mir: Option<bool> = (None, parse_opt_bool, [TRACKED],
         "enable MIR inlining (default: no)"),
+    inline_mir_forwarder_threshold: Option<usize> = (None, parse_opt_number, [TRACKED],
+        "inlining threshold when the caller is a simple forwarding function (default: 30)"),
     inline_mir_hint_threshold: Option<usize> = (None, parse_opt_number, [TRACKED],
         "inlining threshold for functions with inline hint (default: 100)"),
     inline_mir_preserve_debug: Option<bool> = (None, parse_opt_bool, [TRACKED],

library/alloc/src/raw_vec.rs

@@ -429,6 +429,7 @@ impl<T, A: Allocator> RawVec<T, A> {
     ///
     /// Aborts on OOM.
     #[cfg(not(no_global_oom_handling))]
+    #[inline]
     pub fn shrink_to_fit(&mut self, cap: usize) {
         if let Err(err) = self.shrink(cap) {
             handle_error(err);
@@ -511,9 +512,25 @@ impl<T, A: Allocator> RawVec<T, A> {
     }
 
     #[cfg(not(no_global_oom_handling))]
+    #[inline]
     fn shrink(&mut self, cap: usize) -> Result<(), TryReserveError> {
         assert!(cap <= self.capacity(), "Tried to shrink to a larger capacity");
+        // SAFETY: Just checked this isn't trying to grow
+        unsafe { self.shrink_unchecked(cap) }
+    }
 
+    /// `shrink`, but without the capacity check.
+    ///
+    /// This is split out so that `shrink` can inline the check, since it
+    /// optimizes out in things like `shrink_to_fit`, without needing to
+    /// also inline all this code, as doing that ends up failing the
+    /// `vec-shrink-panic` codegen test when `shrink_to_fit` ends up being too
+    /// big for LLVM to be willing to inline.
+    ///
+    /// # Safety
+    /// `cap <= self.capacity()`
+    #[cfg(not(no_global_oom_handling))]
+    unsafe fn shrink_unchecked(&mut self, cap: usize) -> Result<(), TryReserveError> {
         let (ptr, layout) = if let Some(mem) = self.current_memory() { mem } else { return Ok(()) };
         // See current_memory() why this assert is here
         const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };

library/alloc/src/vec/mod.rs

@@ -1101,6 +1101,7 @@ impl<T, A: Allocator> Vec<T, A> {
     /// ```
     #[cfg(not(no_global_oom_handling))]
     #[stable(feature = "rust1", since = "1.0.0")]
+    #[inline]
     pub fn shrink_to_fit(&mut self) {
         // The capacity is never less than the length, and there's nothing to do when
         // they are equal, so we can avoid the panic case in `RawVec::shrink_to_fit`

tests/codegen-units/item-collection/drop_in_place_intrinsic.rs

@@ -1,6 +1,7 @@
 //
 //@ compile-flags:-Zprint-mono-items=eager
 //@ compile-flags:-Zinline-in-all-cgus
+//@ compile-flags:-Zinline-mir=no
 
 #![feature(start)]
 

tests/mir-opt/inline/inline_more_in_non_inline.marked_inline_direct.Inline.after.panic-abort.mir

@@ -0,0 +1,22 @@
+// MIR for `marked_inline_direct` after Inline
+
+fn marked_inline_direct(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        _2 = call_twice(move _3) -> [return: bb1, unwind unreachable];
+    }
+
+    bb1: {
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+}

tests/mir-opt/inline/inline_more_in_non_inline.marked_inline_direct.Inline.after.panic-unwind.mir

@@ -0,0 +1,22 @@
+// MIR for `marked_inline_direct` after Inline
+
+fn marked_inline_direct(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        _2 = call_twice(move _3) -> [return: bb1, unwind continue];
+    }
+
+    bb1: {
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+}

tests/mir-opt/inline/inline_more_in_non_inline.marked_inline_indirect.Inline.after.panic-abort.mir

@@ -0,0 +1,27 @@
+// MIR for `marked_inline_indirect` after Inline
+
+fn marked_inline_indirect(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+    scope 1 (inlined marked_inline_direct) {
+        let _4: ();
+    }
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        StorageLive(_4);
+        _4 = call_twice(move _3) -> [return: bb1, unwind unreachable];
+    }
+
+    bb1: {
+        StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+}

tests/mir-opt/inline/inline_more_in_non_inline.marked_inline_indirect.Inline.after.panic-unwind.mir

@@ -0,0 +1,27 @@
+// MIR for `marked_inline_indirect` after Inline
+
+fn marked_inline_indirect(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+    scope 1 (inlined marked_inline_direct) {
+        let _4: ();
+    }
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        StorageLive(_4);
+        _4 = call_twice(move _3) -> [return: bb1, unwind continue];
+    }
+
+    bb1: {
+        StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+}