Auto merge of rust-lang#127113 - scottmcm:retune-inlining-again, r=<try>

Avoid MIR bloat in inlining

In rust-lang#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.

r? ghost

compiler/rustc_mir_transform/src/cost_checker.rs

@@ -9,6 +9,9 @@ const LANDINGPAD_PENALTY: usize = 50;
 const RESUME_PENALTY: usize = 45;
 const LARGE_SWITCH_PENALTY: usize = 20;
 const CONST_SWITCH_BONUS: usize = 10;
+const MULTIPLE_MUT_PENALTY: usize = 50;
+const REF_COPY_BONUS: usize = 5;
+const MANY_PARAMETERS_BONUS: usize = 10;
 
 /// Verify that the callee body is compatible with the caller.
 #[derive(Clone)]
@@ -31,6 +34,75 @@ 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) {
+        // There's usually extra stack costs to passing lots of parameters,
+        // so we'd rather inline that if the method isn't actually complex,
+        // especially for trait impls that ignore some parameters.
+        if self.callee_body.arg_count > 2 {
+            self.bonus += MANY_PARAMETERS_BONUS;
+        }
+
+        fn is_call_like(bbd: &BasicBlockData<'_>) -> bool {
+            use TerminatorKind::*;
+            match bbd.terminator().kind {
+                Call { .. }
+                | Drop { .. }
+                | Assert { .. }
+                | Yield { .. }
+                | CoroutineDrop
+                | InlineAsm { .. } => true,
+
+                Goto { .. }
+                | SwitchInt { .. }
+                | UnwindResume
+                | UnwindTerminate(_)
+                | Return
+                | Unreachable => false,
+
+                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;
+        }
+
+        let callee_param_env = std::cell::LazyCell::new(|| {
+            let def_id = self.callee_body.source.def_id();
+            self.tcx.param_env(def_id)
+        });
+
+        let mut mut_ref_count = 0;
+        for local in self.callee_body.args_iter() {
+            let ty = self.callee_body.local_decls[local].ty;
+            if let ty::Ref(_, referee, mtbl) = ty.kind() {
+                match mtbl {
+                    Mutability::Mut => mut_ref_count += 1,
+                    Mutability::Not => {
+                        // Encourage inlining `&impl Copy` parameters,
+                        // as that often lets us remove the indirection.
+                        if referee.is_copy_modulo_regions(self.tcx, *callee_param_env) {
+                            self.bonus += REF_COPY_BONUS;
+                        }
+                    }
+                }
+            }
+        }
+
+        // Discourage inlining something with multiple `&mut` parameters,
+        // as the MIR inliner doesn't know how to preserve `noalias`.
+        if mut_ref_count > 1 {
+            self.penalty += MULTIPLE_MUT_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_shim: matches!(
+            codegen_fn_attrs.inline,
+            InlineAttr::Hint | InlineAttr::Always
+        ) && body_is_shim(body),
     };
     let blocks = START_BLOCK..body.basic_blocks.next_index();
     this.process_blocks(body, blocks);
@@ -111,6 +116,7 @@ struct Inliner<'tcx> {
     history: Vec<DefId>,
     /// Indicates that the caller body has been modified.
     changed: bool,
+    caller_is_inline_shim: bool,
 }
 
 impl<'tcx> Inliner<'tcx> {
@@ -485,7 +491,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_shim {
+            self.tcx.sess.opts.unstable_opts.inline_mir_shim_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 +512,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 +1101,37 @@ fn try_instance_mir<'tcx>(
     }
     Ok(tcx.instance_mir(instance))
 }
+
+fn body_is_shim(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

@@ -1771,6 +1771,8 @@ options! {
     inline_mir_preserve_debug: Option<bool> = (None, parse_opt_bool, [TRACKED],
         "when MIR inlining, whether to preserve debug info for callee variables \
         (default: preserve for debuginfo != None, otherwise remove)"),
+    inline_mir_shim_threshold: Option<usize> = (None, parse_opt_number, [TRACKED],
+        "a default MIR inlining threshold (default: 30)"),
     inline_mir_threshold: Option<usize> = (None, parse_opt_number, [TRACKED],
         "a default MIR inlining threshold (default: 50)"),
     input_stats: bool = (false, parse_bool, [UNTRACKED],

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;
+    }
+}

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

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

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

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

tests/mir-opt/inline/inline_more_in_non_inline.rs

@@ -0,0 +1,46 @@
+// EMIT_MIR_FOR_EACH_PANIC_STRATEGY
+//@ compile-flags: -O --crate-type lib
+
+// To avoid MIR blow-up, don't inline large callees into simple shim callers,
+// but *do* inline other trivial things.
+
+extern "Rust" {
+    fn other_thing(x: i32);
+}
+
+#[inline]
+unsafe fn call_twice(x: i32) {
+    unsafe {
+        other_thing(x);
+        other_thing(x);
+    }
+}
+
+// EMIT_MIR inline_more_in_non_inline.monomorphic_not_inline.Inline.after.mir
+#[no_mangle]
+pub unsafe fn monomorphic_not_inline(x: i32) {
+    // CHECK-LABEL: monomorphic_not_inline
+    // CHECK: other_thing
+    // CHECK: other_thing
+    unsafe { call_twice(x) };
+}
+
+// EMIT_MIR inline_more_in_non_inline.marked_inline_direct.Inline.after.mir
+#[inline]
+pub unsafe fn marked_inline_direct(x: i32) {
+    // CHECK-LABEL: marked_inline_direct
+    // CHECK-NOT: other_thing
+    // CHECK: call_twice
+    // CHECK-NOT: other_thing
+    unsafe { call_twice(x) };
+}
+
+// EMIT_MIR inline_more_in_non_inline.marked_inline_indirect.Inline.after.mir
+#[inline]
+pub unsafe fn marked_inline_indirect(x: i32) {
+    // CHECK-LABEL: marked_inline_indirect
+    // CHECK-NOT: other_thing
+    // CHECK: call_twice
+    // CHECK-NOT: other_thing
+    unsafe { marked_inline_direct(x) };
+}