Cranelift: Translate ir::UserExternalNameRefs into callers when inlining (#11389)

This was an entity that we forgot to translate from the callee into the
caller. Note that we do not use the `EntityMap` offset approach for these
entities because `ir::Function` hash-conses them.

Author: fitzgen

cranelift/codegen/src/inline.rs

@@ -132,8 +132,15 @@ pub(crate) fn do_inlining(
                     func_ref,
                 } => {
                     let args = cursor.func.dfg.inst_args(inst);
+                    trace!(
+                        "considering call site for inlining: {inst}: {}",
+                        cursor.func.dfg.display_inst(inst),
+                    );
                     match inliner.inline(&cursor.func, inst, opcode, func_ref, args) {
-                        InlineCommand::KeepCall => continue,
+                        InlineCommand::KeepCall => {
+                            trace!("  --> keeping call");
+                            continue;
+                        }
                         InlineCommand::Inline(callee) => {
                             inline_one(
                                 &mut allocs,
@@ -157,8 +164,15 @@ pub(crate) fn do_inlining(
                     exception,
                 } => {
                     let args = cursor.func.dfg.inst_args(inst);
+                    trace!(
+                        "considering call site for inlining: {inst}: {}",
+                        cursor.func.dfg.display_inst(inst),
+                    );
                     match inliner.inline(&cursor.func, inst, opcode, func_ref, args) {
-                        InlineCommand::KeepCall => continue,
+                        InlineCommand::KeepCall => {
+                            trace!("  --> keeping call");
+                            continue;
+                        }
                         InlineCommand::Inline(callee) => {
                             inline_one(
                                 &mut allocs,
@@ -195,8 +209,18 @@ struct InliningAllocs {
     values: SecondaryMap<ir::Value, PackedOption<ir::Value>>,
 
     /// Map from callee constant to inlined caller constant.
+    ///
+    /// Not in `EntityMap` because these are hash-consed inside the
+    /// `ir::Function`.
     constants: SecondaryMap<ir::Constant, PackedOption<ir::Constant>>,
 
+    /// Map from callee to inlined caller external name refs.
+    ///
+    /// Not in `EntityMap` because these are hash-consed inside the
+    /// `ir::Function`.
+    user_external_name_refs:
+        SecondaryMap<ir::UserExternalNameRef, PackedOption<ir::UserExternalNameRef>>,
+
     /// The set of _caller_ inlined call instructions that need exception table
     /// fixups at the end of inlining.
     ///
@@ -220,6 +244,7 @@ impl InliningAllocs {
         let InliningAllocs {
             values,
             constants,
+            user_external_name_refs,
             calls_needing_exception_table_fixup,
         } = self;
 
@@ -229,6 +254,9 @@ impl InliningAllocs {
         constants.clear();
         constants.resize(callee.dfg.constants.len());
 
+        user_external_name_refs.clear();
+        user_external_name_refs.resize(callee.params.user_named_funcs().len());
+
         // Note: We do not reserve capacity for
         // `calls_needing_exception_table_fixup` because it is a sparse set and
         // we don't know how large it needs to be ahead of time.
@@ -1205,7 +1233,8 @@ fn create_entities(
     entity_map.block_offset = Some(create_blocks(allocs, func, callee));
     entity_map.global_value_offset = Some(create_global_values(func, callee));
     entity_map.sig_ref_offset = Some(create_sig_refs(func, callee));
-    entity_map.func_ref_offset = Some(create_func_refs(func, callee, &entity_map));
+    create_user_external_name_refs(allocs, func, callee);
+    entity_map.func_ref_offset = Some(create_func_refs(allocs, func, callee, &entity_map));
     entity_map.stack_slot_offset = Some(create_stack_slots(func, callee));
     entity_map.dynamic_type_offset = Some(create_dynamic_types(func, callee, &entity_map));
     entity_map.dynamic_stack_slot_offset =
@@ -1309,8 +1338,24 @@ fn create_sig_refs(func: &mut ir::Function, callee: &ir::Function) -> u32 {
     offset
 }
 
+fn create_user_external_name_refs(
+    allocs: &mut InliningAllocs,
+    func: &mut ir::Function,
+    callee: &ir::Function,
+) {
+    for (callee_named_func_ref, name) in callee.params.user_named_funcs().iter() {
+        let caller_named_func_ref = func.declare_imported_user_function(name.clone());
+        allocs.user_external_name_refs[callee_named_func_ref] = Some(caller_named_func_ref).into();
+    }
+}
+
 /// Translate `ir::FuncRef`s from the callee into the caller.
-fn create_func_refs(func: &mut ir::Function, callee: &ir::Function, entity_map: &EntityMap) -> u32 {
+fn create_func_refs(
+    allocs: &InliningAllocs,
+    func: &mut ir::Function,
+    callee: &ir::Function,
+    entity_map: &EntityMap,
+) -> u32 {
     let offset = func.dfg.ext_funcs.len();
     let offset = u32::try_from(offset).unwrap();
 
@@ -1322,7 +1367,17 @@ fn create_func_refs(func: &mut ir::Function, callee: &ir::Function, entity_map:
     } in callee.dfg.ext_funcs.values()
     {
         func.dfg.ext_funcs.push(ir::ExtFuncData {
-            name: name.clone(),
+            name: match name {
+                ir::ExternalName::User(name_ref) => {
+                    ir::ExternalName::User(allocs.user_external_name_refs[*name_ref].expect(
+                        "should have translated all `ir::UserExternalNameRef`s before translating \
+                         `ir::FuncRef`s",
+                    ))
+                }
+                ir::ExternalName::TestCase(_)
+                | ir::ExternalName::LibCall(_)
+                | ir::ExternalName::KnownSymbol(_) => name.clone(),
+            },
             signature: entity_map.inlined_sig_ref(*signature),
             colocated: *colocated,
         });

cranelift/filetests/filetests/inline/user-external-name-refs.clif

@@ -0,0 +1,55 @@
+test inline precise-output
+target x86_64
+
+function %callee() -> i32 {
+    sig0 = () -> i32
+    fn0 = colocated u0:36 sig0
+    fn1 = colocated u0:42 sig0
+block0:
+    v0 = call fn0()
+    v1 = call fn1()
+    v2 = iadd v0, v1
+    return v2
+}
+
+; (no functions inlined into %callee)
+
+function %caller(i32) -> i32 {
+    sig0 = (i32) -> i32
+    sig1 = () -> i32
+    fn0 = colocated u0:1234 sig0
+    fn1 = colocated u0:36 sig0
+    fn2 = %callee sig1
+block0(v0: i32):
+    v1 = call fn0(v0)
+    v2 = call fn2()
+    v3 = call fn1(v1)
+    return v3
+}
+
+; function %caller(i32) -> i32 fast {
+;     sig0 = (i32) -> i32 fast
+;     sig1 = () -> i32 fast
+;     sig2 = () -> i32 fast
+;     fn0 = colocated u0:1234 sig0
+;     fn1 = colocated u0:36 sig0
+;     fn2 = %callee sig1
+;     fn3 = colocated u0:36 sig2
+;     fn4 = colocated u0:42 sig2
+;
+; block0(v0: i32):
+;     v1 = call fn0(v0)
+;     jump block1
+;
+; block1:
+;     v5 = call fn3()
+;     v6 = call fn4()
+;     v7 = iadd v5, v6
+;     jump block2(v7)
+;
+; block2(v4: i32):
+;     v2 -> v4
+;     v3 = call fn1(v1)
+;     return v3
+; }
+

crates/wasmtime/src/compile.rs

@@ -840,6 +840,10 @@ the use case.
                         None
                     }
                 }));
+                log::trace!(
+                    "call graph edges for {output_index:?} = {:?}: {calls:?}",
+                    output.key
+                );
                 Ok(())
             }
         })?;
@@ -865,6 +869,7 @@ the use case.
                 &mut layer_outputs,
                 |output: &mut CompileOutput<'_>| {
                     debug_assert_eq!(output.key.kind(), CompileKind::WasmFunction);
+                    log::trace!("processing inlining for {:?}", output.key);
 
                     let caller_translation = output.translation.unwrap();
                     let caller_module = output.key.module();
@@ -885,11 +890,15 @@ the use case.
                             {
                                 (caller_module, def_func, Some(caller_needs_gc_heap))
                             } else {
-                                let (def_module, def_func) =
-                                    caller_translation.known_imported_functions[callee].expect(
-                                        "a direct call to an imported function must have a \
-                                         statically-known import",
-                                    );
+                                let (def_module, def_func) = caller_translation
+                                    .known_imported_functions[callee]
+                                    .unwrap_or_else(|| {
+                                        panic!(
+                                            "a direct call to an imported function must have a \
+                                             statically-known definition, but direct call to imported \
+                                             function {callee:?} has no statically-known definition",
+                                        )
+                                    });
                                 (def_module, def_func, None)
                             };
 

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

@@ -0,0 +1,96 @@
+;;! target = "x86_64"
+;;! test = "optimize"
+;;! filter = "wasm[2]--function"
+;;! flags = "-C inlining=y"
+
+(component
+  (core module $A
+    (func (export "f0") (result i32)
+      (i32.const 100)
+    )
+    (func (export "f1") (result i32)
+      (i32.const 101)
+    )
+  )
+
+  (core module $B
+    (import "a" "f0" (func $f0 (result i32)))
+    (import "a" "f1" (func $f1 (result i32)))
+    (func (export "f2") (result i32)
+      (i32.add (call $f0) (call $f1))
+    )
+  )
+
+  (core module $C
+    (import "b" "f2" (func $f2 (result i32)))
+    (func (export "f3") (result i32)
+      (i32.add (i32.const 100) (call $f2))
+    )
+  )
+
+  (core instance $a (instantiate $A))
+  (core instance $b (instantiate $B (with "a" (instance $a))))
+  (core instance $c (instantiate $C (with "b" (instance $b))))
+
+  (func (export "f") (result u32)
+    (canon lift (core func $c "f3"))
+  )
+)
+
+;; function u0:1(i64 vmctx, i64) -> i32 tail {
+;;     gv0 = vmctx
+;;     gv1 = load.i64 notrap aligned readonly gv0+8
+;;     gv2 = load.i64 notrap aligned gv1+16
+;;     gv3 = vmctx
+;;     gv4 = vmctx
+;;     gv5 = load.i64 notrap aligned readonly gv4+8
+;;     gv6 = load.i64 notrap aligned gv5+16
+;;     gv7 = vmctx
+;;     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) -> i32 tail
+;;     fn0 = colocated u0:0 sig0
+;;     fn1 = colocated u0:0 sig1
+;;     fn2 = colocated u0:1 sig2
+;;     stack_limit = gv2
+;;
+;;                                 block0(v0: i64, v1: i64):
+;; @00c3                               jump block2
+;;
+;;                                 block2:
+;;                                     jump block4
+;;
+;;                                 block4:
+;;                                     jump block5
+;;
+;;                                 block5:
+;;                                     jump block6
+;;
+;;                                 block6:
+;;                                     jump block7
+;;
+;;                                 block7:
+;;                                     jump block8
+;;
+;;                                 block8:
+;;                                     jump block9
+;;
+;;                                 block9:
+;;                                     jump block3
+;;
+;;                                 block3:
+;;                                     jump block10
+;;
+;;                                 block10:
+;; @00c6                               jump block1
+;;
+;;                                 block1:
+;;                                     v26 = iconst.i32 301
+;; @00c6                               return v26  ; v26 = 301
+;; }