[wasm] Replace the js-to-wasm wrapper eagerly for all matching functions

This CL aims at avoiding compilation of the same js-to-wasm wrapper
multiple times by iterating over all exported functions in the export
table and replacing the wrapper for all functions that share the same
signature with the function that tiered up.

Bug: v8:10982
Change-Id: I721de2f48844349de8a5d12f512a74957c66a0e1
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2527082
Commit-Queue: Vicky Kontoura <vkont@google.com>
Reviewed-by: Jakob Kummerow <jkummerow@chromium.org>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Cr-Commit-Position: refs/heads/master@{#71103}

src/runtime/runtime-wasm.cc

@@ -213,6 +213,20 @@ RUNTIME_FUNCTION(Runtime_WasmCompileLazy) {
   return Object(entrypoint);
 }
 
+namespace {
+void ReplaceWrapper(Isolate* isolate, Handle<WasmInstanceObject> instance,
+                    int function_index, Handle<Code> wrapper_code) {
+  Handle<WasmExternalFunction> exported_function =
+      WasmInstanceObject::GetWasmExternalFunction(isolate, instance,
+                                                  function_index)
+          .ToHandleChecked();
+  exported_function->set_code(*wrapper_code);
+  WasmExportedFunctionData function_data =
+      exported_function->shared().wasm_exported_function_data();
+  function_data.set_wrapper_code(*wrapper_code);
+}
+}  // namespace
+
 RUNTIME_FUNCTION(Runtime_WasmCompileWrapper) {
   HandleScope scope(isolate);
   DCHECK_EQ(2, args.length());
@@ -226,25 +240,40 @@ RUNTIME_FUNCTION(Runtime_WasmCompileWrapper) {
   const wasm::WasmFunction function = module->functions[function_index];
   const wasm::FunctionSig* sig = function.sig;
 
-  MaybeHandle<WasmExternalFunction> maybe_result =
+  // The start function is not guaranteed to be registered as
+  // an exported function (although it is called as one).
+  // If there is no entry for the start function,
+  // the tier-up is abandoned.
+  MaybeHandle<WasmExternalFunction> maybe_exported_function =
       WasmInstanceObject::GetWasmExternalFunction(isolate, instance,
                                                   function_index);
-
-  Handle<WasmExternalFunction> result;
-  if (!maybe_result.ToHandle(&result)) {
-    // We expect the result to be empty in the case of the start function,
-    // which is not an exported function to begin with.
+  Handle<WasmExternalFunction> exported_function;
+  if (!maybe_exported_function.ToHandle(&exported_function)) {
     DCHECK_EQ(function_index, module->start_function_index);
     return ReadOnlyRoots(isolate).undefined_value();
   }
 
-  Handle<Code> wrapper =
+  Handle<Code> wrapper_code =
       wasm::JSToWasmWrapperCompilationUnit::CompileSpecificJSToWasmWrapper(
           isolate, sig, module);
 
-  result->set_code(*wrapper);
+  // Replace the wrapper for the function that triggered the tier-up.
+  // This is to verify that the wrapper is replaced, even if the function
+  // is implicitly exported and is not part of the export_table.
+  ReplaceWrapper(isolate, instance, function_index, wrapper_code);
 
-  function_data->set_wrapper_code(*wrapper);
+  // Iterate over all exports to replace eagerly the wrapper for all functions
+  // that share the signature of the function that tiered up.
+  for (wasm::WasmExport exp : module->export_table) {
+    if (exp.kind != wasm::kExternalFunction) {
+      continue;
+    }
+    int index = static_cast<int>(exp.index);
+    wasm::WasmFunction function = module->functions[index];
+    if (function.sig == sig && index != function_index) {
+      ReplaceWrapper(isolate, instance, index, wrapper_code);
+    }
+  }
 
   return ReadOnlyRoots(isolate).undefined_value();
 }

test/cctest/wasm/test-run-wasm-wrappers.cc

@@ -176,6 +176,144 @@ TEST(WrapperReplacement) {
   }
   Cleanup();
 }
+
+TEST(EagerWrapperReplacement) {
+  {
+    // This test assumes use of the generic wrapper.
+    FlagScope<bool> use_wasm_generic_wrapper(&FLAG_wasm_generic_wrapper, true);
+
+    TestSignatures sigs;
+    AccountingAllocator allocator;
+    Zone zone(&allocator, ZONE_NAME);
+
+    // Define three Wasm functions.
+    // Two of these functions (add and mult) will share the same signature,
+    // while the other one (id) won't.
+    WasmModuleBuilder* builder = zone.New<WasmModuleBuilder>(&zone);
+    WasmFunctionBuilder* add = builder->AddFunction(sigs.i_ii());
+    add->builder()->AddExport(CStrVector("add"), add);
+    byte add_code[] = {WASM_I32_ADD(WASM_GET_LOCAL(0), WASM_GET_LOCAL(1)),
+                       WASM_END};
+    add->EmitCode(add_code, sizeof(add_code));
+    WasmFunctionBuilder* mult = builder->AddFunction(sigs.i_ii());
+    mult->builder()->AddExport(CStrVector("mult"), mult);
+    byte mult_code[] = {WASM_I32_MUL(WASM_GET_LOCAL(0), WASM_GET_LOCAL(1)),
+                        WASM_END};
+    mult->EmitCode(mult_code, sizeof(mult_code));
+    WasmFunctionBuilder* id = builder->AddFunction(sigs.i_i());
+    id->builder()->AddExport(CStrVector("id"), id);
+    byte id_code[] = {WASM_GET_LOCAL(0), WASM_END};
+    id->EmitCode(id_code, sizeof(id_code));
+
+    // Compile module.
+    ZoneBuffer buffer(&zone);
+    builder->WriteTo(&buffer);
+    Isolate* isolate = CcTest::InitIsolateOnce();
+    HandleScope scope(isolate);
+    testing::SetupIsolateForWasmModule(isolate);
+    ErrorThrower thrower(isolate, "CompileAndRunWasmModule");
+    MaybeHandle<WasmInstanceObject> maybe_instance =
+        CompileAndInstantiateForTesting(
+            isolate, &thrower, ModuleWireBytes(buffer.begin(), buffer.end()));
+    Handle<WasmInstanceObject> instance = maybe_instance.ToHandleChecked();
+
+    // Get the exported functions.
+    Handle<WasmExportedFunction> add_export =
+        testing::GetExportedFunction(isolate, instance, "add")
+            .ToHandleChecked();
+    Handle<WasmExportedFunction> mult_export =
+        testing::GetExportedFunction(isolate, instance, "mult")
+            .ToHandleChecked();
+    Handle<WasmExportedFunction> id_export =
+        testing::GetExportedFunction(isolate, instance, "id").ToHandleChecked();
+
+    // Get the function data for all exported functions.
+    WasmExportedFunctionData add_function_data =
+        add_export->shared().wasm_exported_function_data();
+    WasmExportedFunctionData mult_function_data =
+        mult_export->shared().wasm_exported_function_data();
+    WasmExportedFunctionData id_function_data =
+        id_export->shared().wasm_exported_function_data();
+
+    // Set the call count for add to (threshold - 1),
+    // so that the next call to it will cause the function to tier up.
+    add_function_data.set_call_count(kGenericWrapperThreshold - 1);
+
+    // Verify that the call counts for all functions are correct.
+    CHECK_EQ(add_function_data.call_count(), kGenericWrapperThreshold - 1);
+    CHECK_EQ(mult_function_data.call_count(), 0);
+    CHECK_EQ(id_function_data.call_count(), 0);
+
+    // Verify that all functions are set to use the generic wrapper.
+    CHECK(add_function_data.wrapper_code().is_builtin() &&
+          add_function_data.wrapper_code().builtin_index() ==
+              Builtins::kGenericJSToWasmWrapper);
+    CHECK(mult_function_data.wrapper_code().is_builtin() &&
+          mult_function_data.wrapper_code().builtin_index() ==
+              Builtins::kGenericJSToWasmWrapper);
+    CHECK(id_function_data.wrapper_code().is_builtin() &&
+          id_function_data.wrapper_code().builtin_index() ==
+              Builtins::kGenericJSToWasmWrapper);
+
+    // Call the add function to trigger the tier up.
+    {
+      int32_t expected_value = 21;
+      Handle<Object> params[2] = {Handle<Object>(Smi::FromInt(10), isolate),
+                                  Handle<Object>(Smi::FromInt(11), isolate)};
+      Handle<Object> receiver = isolate->factory()->undefined_value();
+      MaybeHandle<Object> maybe_result =
+          Execution::Call(isolate, add_export, receiver, 2, params);
+      Handle<Object> result = maybe_result.ToHandleChecked();
+      CHECK(result->IsSmi() && Smi::ToInt(*result) == expected_value);
+    }
+
+    // Verify that the call counts for all functions are correct.
+    CHECK_EQ(add_function_data.call_count(), kGenericWrapperThreshold);
+    CHECK_EQ(mult_function_data.call_count(), 0);
+    CHECK_EQ(id_function_data.call_count(), 0);
+
+    // Verify that the tier up of the add function replaced the wrapper
+    // for both the add and the mult functions, but not the id function.
+    CHECK(add_function_data.wrapper_code().kind() ==
+          CodeKind::JS_TO_WASM_FUNCTION);
+    CHECK(mult_function_data.wrapper_code().kind() ==
+          CodeKind::JS_TO_WASM_FUNCTION);
+    CHECK(id_function_data.wrapper_code().is_builtin() &&
+          id_function_data.wrapper_code().builtin_index() ==
+              Builtins::kGenericJSToWasmWrapper);
+
+    // Call the mult function to verify that the compiled wrapper is used.
+    {
+      int32_t expected_value = 42;
+      Handle<Object> params[2] = {Handle<Object>(Smi::FromInt(7), isolate),
+                                  Handle<Object>(Smi::FromInt(6), isolate)};
+      Handle<Object> receiver = isolate->factory()->undefined_value();
+      MaybeHandle<Object> maybe_result =
+          Execution::Call(isolate, mult_export, receiver, 2, params);
+      Handle<Object> result = maybe_result.ToHandleChecked();
+      CHECK(result->IsSmi() && Smi::ToInt(*result) == expected_value);
+    }
+    // Verify that mult's call count is still 0, which means that the call
+    // didn't go through the generic wrapper.
+    CHECK_EQ(mult_function_data.call_count(), 0);
+
+    // Call the id function to verify that the generic wrapper is used.
+    {
+      int32_t expected_value = 12;
+      Handle<Object> params[1] = {
+          Handle<Object>(Smi::FromInt(expected_value), isolate)};
+      Handle<Object> receiver = isolate->factory()->undefined_value();
+      MaybeHandle<Object> maybe_result =
+          Execution::Call(isolate, id_export, receiver, 1, params);
+      Handle<Object> result = maybe_result.ToHandleChecked();
+      CHECK(result->IsSmi() && Smi::ToInt(*result) == expected_value);
+    }
+    // Verify that id's call count increased to 1, which means that the call
+    // used the generic wrapper.
+    CHECK_EQ(id_function_data.call_count(), 1);
+  }
+  Cleanup();
+}
 #endif
 
 }  // namespace test_run_wasm_wrappers