Do not consider inlined function bodies for more inlining (#11462)

Due to our bottom-up, SCC-based approach to inlining, we've already performed
inlining on the callee, and do not need to consider any further inlining within
it.

This also prevents us from inlining recursive callees all the way up to the size
threshold, and instead just inline one layer.

Author: fitzgen

cranelift/codegen/src/inline.rs

@@ -37,12 +37,19 @@ type SmallBlockCallVec = SmallVec<[ir::BlockCall; 8]>;
 pub enum InlineCommand<'a> {
     /// Keep the call as-is, out-of-line, and do not inline the callee.
     KeepCall,
+
     /// Inline the call, using this function as the body of the callee.
     ///
     /// It is the `Inline` implementor's responsibility to ensure that this
     /// function is the correct callee. Providing the wrong function may result
     /// in panics during compilation or incorrect runtime behavior.
-    Inline(Cow<'a, ir::Function>),
+    Inline {
+        /// The callee function's body.
+        callee: Cow<'a, ir::Function>,
+        /// Whether to visit any function calls within the callee body after
+        /// inlining and consider them for further inlining.
+        visit_callee: bool,
+    },
 }
 
 /// A trait for directing Cranelift whether to inline a particular call or not.
@@ -112,37 +119,45 @@ pub(crate) fn do_inlining(
     let mut allocs = InliningAllocs::default();
 
     let mut cursor = FuncCursor::new(func);
-    while let Some(block) = cursor.next_block() {
-        // Always keep track of our previous cursor position. After we inline a
-        // call, replacing the current position with an arbitrary sub-CFG, we
-        // back up to this previous position. This makes sure our cursor is
-        // always at a position that is inserted in the layout and also enables
-        // multi-level inlining, if desired by the user, where we consider any
-        // newly-inlined call instructions for further inlining.
+    'block_loop: while let Some(block) = cursor.next_block() {
+        // Always keep track of our previous cursor position. Assuming that the
+        // current position is a function call that we will inline, then the
+        // previous position is just before the inlined callee function. After
+        // inlining a call, the Cranelift user can decide whether to consider
+        // any function calls in the inlined callee for further inlining or
+        // not. When they do, then we back up to this previous cursor position
+        // so that our traversal will then continue over the inlined body.
         let mut prev_pos;
 
         while let Some(inst) = {
             prev_pos = cursor.position();
             cursor.next_inst()
         } {
+            // Make sure that `block` is always `inst`'s block, even with all of
+            // our cursor-position-updating and block-splitting-during-inlining
+            // shenanigans below.
+            debug_assert_eq!(Some(block), cursor.func.layout.inst_block(inst));
+
             match cursor.func.dfg.insts[inst] {
                 ir::InstructionData::Call {
                     opcode: opcode @ ir::Opcode::Call | opcode @ ir::Opcode::ReturnCall,
                     args: _,
                     func_ref,
                 } => {
-                    let args = cursor.func.dfg.inst_args(inst);
                     trace!(
                         "considering call site for inlining: {inst}: {}",
                         cursor.func.dfg.display_inst(inst),
                     );
+                    let args = cursor.func.dfg.inst_args(inst);
                     match inliner.inline(&cursor.func, inst, opcode, func_ref, args) {
                         InlineCommand::KeepCall => {
                             trace!("  --> keeping call");
-                            continue;
                         }
-                        InlineCommand::Inline(callee) => {
-                            inline_one(
+                        InlineCommand::Inline {
+                            callee,
+                            visit_callee,
+                        } => {
+                            let last_inlined_block = inline_one(
                                 &mut allocs,
                                 cursor.func,
                                 func_ref,
@@ -153,7 +168,15 @@ pub(crate) fn do_inlining(
                                 None,
                             );
                             inlined_any = true;
-                            cursor.set_position(prev_pos);
+                            if visit_callee {
+                                cursor.set_position(prev_pos);
+                            } else {
+                                // Arrange it so that the `next_block()` loop
+                                // will continue to the next block that is not
+                                // associated with the just-inlined callee.
+                                cursor.goto_bottom(last_inlined_block);
+                                continue 'block_loop;
+                            }
                         }
                     }
                 }
@@ -163,18 +186,20 @@ pub(crate) fn do_inlining(
                     func_ref,
                     exception,
                 } => {
-                    let args = cursor.func.dfg.inst_args(inst);
                     trace!(
                         "considering call site for inlining: {inst}: {}",
                         cursor.func.dfg.display_inst(inst),
                     );
+                    let args = cursor.func.dfg.inst_args(inst);
                     match inliner.inline(&cursor.func, inst, opcode, func_ref, args) {
                         InlineCommand::KeepCall => {
                             trace!("  --> keeping call");
-                            continue;
                         }
-                        InlineCommand::Inline(callee) => {
-                            inline_one(
+                        InlineCommand::Inline {
+                            callee,
+                            visit_callee,
+                        } => {
+                            let last_inlined_block = inline_one(
                                 &mut allocs,
                                 cursor.func,
                                 func_ref,
@@ -185,11 +210,30 @@ pub(crate) fn do_inlining(
                                 Some(exception),
                             );
                             inlined_any = true;
-                            cursor.set_position(prev_pos);
+                            if visit_callee {
+                                cursor.set_position(prev_pos);
+                            } else {
+                                // Arrange it so that the `next_block()` loop
+                                // will continue to the next block that is not
+                                // associated with the just-inlined callee.
+                                cursor.goto_bottom(last_inlined_block);
+                                continue 'block_loop;
+                            }
                         }
                     }
                 }
-                _ => continue,
+                ir::InstructionData::CallIndirect { .. }
+                | ir::InstructionData::TryCallIndirect { .. } => {
+                    // Can't inline indirect calls; need to have some earlier
+                    // pass rewrite them into direct calls first, when possible.
+                }
+                _ => {
+                    debug_assert!(
+                        !cursor.func.dfg.insts[inst].opcode().is_call(),
+                        "should have matched all call instructions, but found: {inst}: {}",
+                        cursor.func.dfg.display_inst(inst),
+                    );
+                }
             }
         }
     }
@@ -283,6 +327,8 @@ impl InliningAllocs {
 }
 
 /// Inline one particular function call.
+///
+/// Returns the last inlined block in the layout.
 fn inline_one(
     allocs: &mut InliningAllocs,
     func: &mut ir::Function,
@@ -292,7 +338,7 @@ fn inline_one(
     call_opcode: ir::Opcode,
     callee: &ir::Function,
     call_exception_table: Option<ir::ExceptionTable>,
-) {
+) -> ir::Block {
     trace!(
         "Inlining call {call_inst:?}: {}\n\
          with callee = {callee:?}",
@@ -318,7 +364,7 @@ fn inline_one(
     // Prepare for translating the actual instructions by inserting the inlined
     // blocks into the caller's layout in the same order that they appear in the
     // callee.
-    inline_block_layout(func, call_block, callee, &entity_map);
+    let last_inlined_block = inline_block_layout(func, call_block, callee, &entity_map);
 
     // Translate each instruction from the callee into the caller,
     // appending them to their associated block in the caller.
@@ -470,6 +516,8 @@ fn inline_one(
     if let Some(call_exception_table) = call_exception_table {
         fixup_inlined_call_exception_tables(allocs, func, call_exception_table);
     }
+
+    last_inlined_block
 }
 
 /// Append stack map entries from the caller and callee to the given inlined
@@ -916,12 +964,14 @@ impl<'a> ir::instructions::InstructionMapper for InliningInstRemapper<'a> {
 }
 
 /// Inline the callee's layout into the caller's layout.
+///
+/// Returns the last inlined block in the layout.
 fn inline_block_layout(
     func: &mut ir::Function,
     call_block: ir::Block,
     callee: &ir::Function,
     entity_map: &EntityMap,
-) {
+) -> ir::Block {
     // Iterate over callee blocks in layout order, inserting their associated
     // inlined block into the caller's layout.
     let mut prev_inlined_block = call_block;
@@ -934,6 +984,7 @@ fn inline_block_layout(
         prev_inlined_block = inlined_block;
         next_callee_block = callee.layout.next_block(callee_block);
     }
+    prev_inlined_block
 }
 
 /// Split the call instruction's block just after the call instruction to create

cranelift/filetests/src/test_inline.rs

@@ -136,12 +136,18 @@ impl<'a> Inline for Inliner<'a> {
                 .and_then(|name| self.0.get(&ir::UserFuncName::User(name.clone())))
             {
                 None => InlineCommand::KeepCall,
-                Some(f) => InlineCommand::Inline(Cow::Borrowed(f)),
+                Some(f) => InlineCommand::Inline {
+                    callee: Cow::Borrowed(f),
+                    visit_callee: true,
+                },
             },
             ir::ExternalName::TestCase(name) => {
                 match self.0.get(&ir::UserFuncName::Testcase(name.clone())) {
                     None => InlineCommand::KeepCall,
-                    Some(f) => InlineCommand::Inline(Cow::Borrowed(f)),
+                    Some(f) => InlineCommand::Inline {
+                        callee: Cow::Borrowed(f),
+                        visit_callee: true,
+                    },
                 }
             }
             ir::ExternalName::LibCall(_) | ir::ExternalName::KnownSymbol(_) => {

crates/cranelift/src/compiler.rs

@@ -845,9 +845,13 @@ impl InliningCompiler for Compiler {
                                 .code
                                 .downcast_ref::<Option<CompilerContext>>()
                                 .unwrap();
-                            InlineCommand::Inline(Cow::Borrowed(
-                                &cx.as_ref().unwrap().codegen_context.func,
-                            ))
+                            InlineCommand::Inline {
+                                callee: Cow::Borrowed(&cx.as_ref().unwrap().codegen_context.func),
+                                // We've already visited the callee for inlining
+                                // due to our bottom-up approach, no need to
+                                // visit it again.
+                                visit_callee: false,
+                            }
                         }
                     },
                     _ => InlineCommand::KeepCall,

tests/disas/component-model/inlining-bug.wat

@@ -46,9 +46,6 @@
 ;;     gv8 = vmctx
 ;;     gv9 = load.i64 notrap aligned readonly gv8+8
 ;;     gv10 = load.i64 notrap aligned gv9+16
-;;     gv11 = vmctx
-;;     gv12 = load.i64 notrap aligned readonly gv11+8
-;;     gv13 = load.i64 notrap aligned gv12+16
 ;;     sig0 = (i64 vmctx, i64) -> i32 tail
 ;;     sig1 = (i64 vmctx, i64) -> i32 tail
 ;;     sig2 = (i64 vmctx, i64) tail
@@ -64,15 +61,9 @@
 ;;                                     jump block4
 ;;
 ;;                                 block4:
-;;                                     jump block7
-;;
-;;                                 block7:
-;;                                     jump block8
-;;
-;;                                 block8:
-;;                                     jump block9
-;;
-;;                                 block9:
+;; @00d4                               v4 = load.i64 notrap aligned readonly can_move v0+64
+;;                                     v10 = load.i64 notrap aligned readonly can_move v4+88
+;;                                     call fn2(v10, v10)
 ;;                                     jump block5
 ;;
 ;;                                 block5:
@@ -82,9 +73,9 @@
 ;;                                     jump block3
 ;;
 ;;                                 block3:
-;;                                     jump block10
+;;                                     jump block7
 ;;
-;;                                 block10:
+;;                                 block7:
 ;; @00d6                               jump block1
 ;;
 ;;                                 block1: