[wasm] Replace generic js-to-wasm wrapper when threshold is reached

This CL adds a basic tiering strategy for the js-to-wasm wrappers.
When applicable, calls to exported WebAssembly functions are initially
handled through the generic js-to-wasm wrapper. If these calls
through the generic wrapper reach a constant threshold, the specific
(per-signature) wrapper is compiled synchronously for the function
and the generic wrapper is replaced.

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

src/builtins/x64/builtins-x64.cc

@@ -26,6 +26,7 @@
 #include "src/objects/smi.h"
 #include "src/wasm/baseline/liftoff-assembler-defs.h"
 #include "src/wasm/object-access.h"
+#include "src/wasm/wasm-constants.h"
 #include "src/wasm/wasm-linkage.h"
 #include "src/wasm/wasm-objects.h"
 
@@ -3056,6 +3057,27 @@ void Builtins::Generate_GenericJSToWasmWrapper(MacroAssembler* masm) {
   // Set up the stackframe.
   __ EnterFrame(StackFrame::JS_TO_WASM);
 
+  // -------------------------------------------
+  // Compute offsets and prepare for GC.
+  // -------------------------------------------
+  // We will have to save a value indicating the GC the number
+  // of values on the top of the stack that have to be scanned before calling
+  // the Wasm function.
+  constexpr int kFrameMarkerOffset = -kSystemPointerSize;
+  constexpr int kGCScanSlotCountOffset =
+      kFrameMarkerOffset - kSystemPointerSize;
+  constexpr int kParamCountOffset = kGCScanSlotCountOffset - kSystemPointerSize;
+  constexpr int kReturnCountOffset = kParamCountOffset - kSystemPointerSize;
+  constexpr int kValueTypesArrayStartOffset =
+      kReturnCountOffset - kSystemPointerSize;
+  // We set and use this slot only when moving parameters into the parameter
+  // registers (so no GC scan is needed).
+  constexpr int kFunctionDataOffset =
+      kValueTypesArrayStartOffset - kSystemPointerSize;
+  constexpr int kLastSpillOffset = kFunctionDataOffset;
+  constexpr int kNumSpillSlots = 5;
+  __ subq(rsp, Immediate(kNumSpillSlots * kSystemPointerSize));
+
   // -------------------------------------------
   // Load the Wasm exported function data and the Wasm instance.
   // -------------------------------------------
@@ -3088,6 +3110,17 @@ void Builtins::Generate_GenericJSToWasmWrapper(MacroAssembler* masm) {
                  WasmExportedFunctionData::kCallCountOffset - kHeapObjectTag),
       Smi::FromInt(1));
 
+  // -------------------------------------------
+  // Check if the call count reached the threshold.
+  // -------------------------------------------
+  Label compile_wrapper, compile_wrapper_done;
+  __ SmiCompare(
+      MemOperand(function_data,
+                 WasmExportedFunctionData::kCallCountOffset - kHeapObjectTag),
+      Smi::FromInt(wasm::kGenericWrapperThreshold));
+  __ j(greater_equal, &compile_wrapper);
+  __ bind(&compile_wrapper_done);
+
   // -------------------------------------------
   // Load values from the signature.
   // -------------------------------------------
@@ -3114,29 +3147,12 @@ void Builtins::Generate_GenericJSToWasmWrapper(MacroAssembler* masm) {
   signature = no_reg;
 
   // -------------------------------------------
-  // Set up the stack.
+  // Store signature-related values to the stack.
   // -------------------------------------------
   // We store values on the stack to restore them after function calls.
   // We cannot push values onto the stack right before the wasm call. The wasm
   // function expects the parameters, that didn't fit into the registers, on the
   // top of the stack.
-  // We will have to save a value indicating the GC the number
-  // of values on the top of the stack that have to be scanned before calling
-  // the Wasm function.
-  constexpr int kFrameMarkerOffset = -kSystemPointerSize;
-  constexpr int kGCScanSlotCountOffset =
-      kFrameMarkerOffset - kSystemPointerSize;
-  constexpr int kParamCountOffset = kGCScanSlotCountOffset - kSystemPointerSize;
-  constexpr int kReturnCountOffset = kParamCountOffset - kSystemPointerSize;
-  constexpr int kValueTypesArrayStartOffset =
-      kReturnCountOffset - kSystemPointerSize;
-  // We set and use this slot only when moving parameters into the parameter
-  // registers (so no GC scan is needed).
-  constexpr int kFunctionDataOffset =
-      kValueTypesArrayStartOffset - kSystemPointerSize;
-  constexpr int kLastSpillOffset = kFunctionDataOffset;
-  constexpr int kNumSpillSlots = 5;
-  __ subq(rsp, Immediate(kNumSpillSlots * kSystemPointerSize));
   __ movq(MemOperand(rbp, kParamCountOffset), param_count);
   __ movq(MemOperand(rbp, kReturnCountOffset), return_count);
   __ movq(MemOperand(rbp, kValueTypesArrayStartOffset), valuetypes_array_ptr);
@@ -3694,6 +3710,30 @@ void Builtins::Generate_GenericJSToWasmWrapper(MacroAssembler* masm) {
   __ Call(BUILTIN_CODE(masm->isolate(), WasmFloat64ToNumber),
           RelocInfo::CODE_TARGET);
   __ jmp(&return_done);
+
+  // -------------------------------------------
+  // Kick off compilation.
+  // -------------------------------------------
+  __ bind(&compile_wrapper);
+  // Enable GC.
+  MemOperand GCScanSlotPlace = MemOperand(rbp, kGCScanSlotCountOffset);
+  __ movq(GCScanSlotPlace, Immediate(4));
+  // Save registers to the stack.
+  __ pushq(wasm_instance);
+  __ pushq(function_data);
+  // Push the arguments for the runtime call.
+  __ Push(wasm_instance);  // first argument
+  __ Push(function_data);  // second argument
+  // Set up context.
+  __ Move(kContextRegister, Smi::zero());
+  // Call the runtime function that kicks off compilation.
+  __ CallRuntime(Runtime::kWasmCompileWrapper, 2);
+  // Pop the result.
+  __ movq(r9, kReturnRegister0);
+  // Restore registers from the stack.
+  __ popq(function_data);
+  __ popq(wasm_instance);
+  __ jmp(&compile_wrapper_done);
 }
 
 namespace {

src/runtime/runtime-wasm.cc

@@ -213,6 +213,34 @@ RUNTIME_FUNCTION(Runtime_WasmCompileLazy) {
   return Object(entrypoint);
 }
 
+RUNTIME_FUNCTION(Runtime_WasmCompileWrapper) {
+  HandleScope scope(isolate);
+  DCHECK_EQ(2, args.length());
+  CONVERT_ARG_HANDLE_CHECKED(WasmInstanceObject, instance, 0);
+  CONVERT_ARG_HANDLE_CHECKED(WasmExportedFunctionData, function_data, 1);
+  DCHECK(isolate->context().is_null());
+  isolate->set_context(instance->native_context());
+
+  const wasm::WasmModule* module = instance->module();
+  const int function_index = function_data->function_index();
+  const wasm::WasmFunction function = module->functions[function_index];
+  const wasm::FunctionSig* sig = function.sig;
+
+  Handle<Code> wrapper =
+      wasm::JSToWasmWrapperCompilationUnit::CompileSpecificJSToWasmWrapper(
+          isolate, sig, module);
+
+  function_data->set_wrapper_code(*wrapper);
+
+  MaybeHandle<WasmExternalFunction> maybe_result =
+      WasmInstanceObject::GetWasmExternalFunction(isolate, instance,
+                                                  function_index);
+  Handle<WasmExternalFunction> result = maybe_result.ToHandleChecked();
+  result->set_code(*wrapper);
+
+  return ReadOnlyRoots(isolate).undefined_value();
+}
+
 RUNTIME_FUNCTION(Runtime_WasmTriggerTierUp) {
   HandleScope scope(isolate);
   DCHECK_EQ(1, args.length());

src/runtime/runtime.h

@@ -579,6 +579,7 @@ namespace internal {
   F(WasmTableFill, 4, 1)              \
   F(WasmIsValidRefValue, 3, 1)        \
   F(WasmCompileLazy, 2, 1)            \
+  F(WasmCompileWrapper, 2, 1)         \
   F(WasmTriggerTierUp, 1, 1)          \
   F(WasmDebugBreak, 0, 1)             \
   F(WasmAllocateRtt, 2, 1)

src/wasm/function-compiler.cc

@@ -294,10 +294,11 @@ bool UseGenericWrapper(const FunctionSig* sig) {
 JSToWasmWrapperCompilationUnit::JSToWasmWrapperCompilationUnit(
     Isolate* isolate, WasmEngine* wasm_engine, const FunctionSig* sig,
     const WasmModule* module, bool is_import,
-    const WasmFeatures& enabled_features)
+    const WasmFeatures& enabled_features, AllowGeneric allow_generic)
     : is_import_(is_import),
       sig_(sig),
-      use_generic_wrapper_(UseGenericWrapper(sig) && !is_import),
+      use_generic_wrapper_(allow_generic && UseGenericWrapper(sig) &&
+                           !is_import),
       job_(use_generic_wrapper_ ? nullptr
                                 : compiler::NewJSToWasmCompilationJob(
                                       isolate, wasm_engine, sig, module,
@@ -338,7 +339,21 @@ Handle<Code> JSToWasmWrapperCompilationUnit::CompileJSToWasmWrapper(
   // Run the compilation unit synchronously.
   WasmFeatures enabled_features = WasmFeatures::FromIsolate(isolate);
   JSToWasmWrapperCompilationUnit unit(isolate, isolate->wasm_engine(), sig,
-                                      module, is_import, enabled_features);
+                                      module, is_import, enabled_features,
+                                      kAllowGeneric);
+  unit.Execute();
+  return unit.Finalize(isolate);
+}
+
+// static
+Handle<Code> JSToWasmWrapperCompilationUnit::CompileSpecificJSToWasmWrapper(
+    Isolate* isolate, const FunctionSig* sig, const WasmModule* module) {
+  // Run the compilation unit synchronously.
+  const bool is_import = false;
+  WasmFeatures enabled_features = WasmFeatures::FromIsolate(isolate);
+  JSToWasmWrapperCompilationUnit unit(isolate, isolate->wasm_engine(), sig,
+                                      module, is_import, enabled_features,
+                                      kDontAllowGeneric);
   unit.Execute();
   return unit.Finalize(isolate);
 }

src/wasm/function-compiler.h

@@ -113,10 +113,15 @@ STATIC_ASSERT(sizeof(WasmCompilationUnit) <= 2 * kSystemPointerSize);
 
 class V8_EXPORT_PRIVATE JSToWasmWrapperCompilationUnit final {
  public:
+  // A flag to mark whether the compilation unit can skip the compilation
+  // and return the builtin (generic) wrapper, when available.
+  enum AllowGeneric : bool { kAllowGeneric = true, kDontAllowGeneric = false };
+
   JSToWasmWrapperCompilationUnit(Isolate* isolate, WasmEngine* wasm_engine,
                                  const FunctionSig* sig,
                                  const wasm::WasmModule* module, bool is_import,
-                                 const WasmFeatures& enabled_features);
+                                 const WasmFeatures& enabled_features,
+                                 AllowGeneric allow_generic);
   ~JSToWasmWrapperCompilationUnit();
 
   void Execute();
@@ -131,6 +136,12 @@ class V8_EXPORT_PRIVATE JSToWasmWrapperCompilationUnit final {
                                              const WasmModule* module,
                                              bool is_import);
 
+  // Run a compilation unit synchronously, but ask for the specific
+  // wrapper.
+  static Handle<Code> CompileSpecificJSToWasmWrapper(Isolate* isolate,
+                                                     const FunctionSig* sig,
+                                                     const WasmModule* module);
+
  private:
   bool is_import_;
   const FunctionSig* sig_;

src/wasm/module-compiler.cc

@@ -1349,7 +1349,8 @@ int AddExportWrapperUnits(Isolate* isolate, WasmEngine* wasm_engine,
     if (keys.insert(key).second) {
       auto unit = std::make_shared<JSToWasmWrapperCompilationUnit>(
           isolate, wasm_engine, function.sig, native_module->module(),
-          function.imported, enabled_features);
+          function.imported, enabled_features,
+          JSToWasmWrapperCompilationUnit::kAllowGeneric);
       builder->AddJSToWasmWrapperUnit(std::move(unit));
     }
   }
@@ -3293,7 +3294,8 @@ void CompileJsToWasmWrappers(Isolate* isolate, const WasmModule* module,
     if (queue.insert(key)) {
       auto unit = std::make_unique<JSToWasmWrapperCompilationUnit>(
           isolate, isolate->wasm_engine(), function.sig, module,
-          function.imported, enabled_features);
+          function.imported, enabled_features,
+          JSToWasmWrapperCompilationUnit::kAllowGeneric);
       compilation_units.emplace(key, std::move(unit));
     }
   }

src/wasm/wasm-constants.h

@@ -123,6 +123,11 @@ constexpr uint32_t kExceptionAttribute = 0;
 
 constexpr int kAnonymousFuncIndex = -1;
 
+// The number of calls to an exported wasm function that will be handled
+// by the generic wrapper. Once this threshold is reached, a specific wrapper
+// is to be compiled for the function's signature.
+constexpr uint32_t kGenericWrapperThreshold = 6;
+
 }  // namespace wasm
 }  // namespace internal
 }  // namespace v8

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

@@ -82,6 +82,100 @@ TEST(CallCounter) {
   }
   Cleanup();
 }
+
+TEST(WrapperReplacement) {
+  {
+    // 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 the Wasm function.
+    WasmModuleBuilder* builder = zone.New<WasmModuleBuilder>(&zone);
+    WasmFunctionBuilder* f = builder->AddFunction(sigs.i_i());
+    f->builder()->AddExport(CStrVector("main"), f);
+    byte code[] = {WASM_RETURN1(WASM_GET_LOCAL(0)), WASM_END};
+    f->EmitCode(code, sizeof(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> instance = CompileAndInstantiateForTesting(
+        isolate, &thrower, ModuleWireBytes(buffer.begin(), buffer.end()));
+
+    // Get the exported function.
+    MaybeHandle<WasmExportedFunction> maybe_export =
+        testing::GetExportedFunction(isolate, instance.ToHandleChecked(),
+                                     "main");
+    Handle<WasmExportedFunction> main_export = maybe_export.ToHandleChecked();
+
+    // Check that the counter has initially a value of 0.
+    CHECK_EQ(main_export->shared().wasm_exported_function_data().call_count(),
+             0);
+    CHECK_GT(kGenericWrapperThreshold, 0);
+
+    // Call the exported Wasm function as many times as required to reach the
+    // threshold for compiling the specific wrapper.
+    const int threshold = static_cast<int>(kGenericWrapperThreshold);
+    for (int i = 1; i < threshold; ++i) {
+      // Verify that the wrapper to be used is still the generic one.
+      Code wrapper =
+          main_export->shared().wasm_exported_function_data().wrapper_code();
+      CHECK(wrapper.is_builtin() &&
+            wrapper.builtin_index() == Builtins::kGenericJSToWasmWrapper);
+      // Call the function.
+      int32_t expected_value = i;
+      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, main_export, receiver, 1, params);
+      Handle<Object> result = maybe_result.ToHandleChecked();
+      // Verify that the counter has now a value of i and the return value is
+      // correct.
+      CHECK_EQ(main_export->shared().wasm_exported_function_data().call_count(),
+               i);
+      CHECK(result->IsSmi() && Smi::ToInt(*result) == expected_value);
+    }
+
+    // Get the wrapper-code object before making the call that will kick off the
+    // wrapper replacement.
+    Code wrapper_before_call =
+        main_export->shared().wasm_exported_function_data().wrapper_code();
+    // Verify that the wrapper before the call is the generic wrapper.
+    CHECK(wrapper_before_call.is_builtin() &&
+          wrapper_before_call.builtin_index() ==
+              Builtins::kGenericJSToWasmWrapper);
+
+    // Call the exported Wasm function one more time to kick off the wrapper
+    // replacement.
+    int32_t expected_value = 42;
+    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, main_export, receiver, 1, params);
+    Handle<Object> result = maybe_result.ToHandleChecked();
+    // Check that the counter has the threshold value and the result is correct.
+    CHECK_EQ(main_export->shared().wasm_exported_function_data().call_count(),
+             kGenericWrapperThreshold);
+    CHECK(result->IsSmi() && Smi::ToInt(*result) == expected_value);
+
+    // Verify that the wrapper-code object has changed.
+    Code wrapper_after_call =
+        main_export->shared().wasm_exported_function_data().wrapper_code();
+    CHECK_NE(wrapper_after_call, wrapper_before_call);
+    // Verify that the wrapper is now a specific one.
+    CHECK(wrapper_after_call.kind() == CodeKind::JS_TO_WASM_FUNCTION);
+  }
+  Cleanup();
+}
 #endif
 
 }  // namespace test_run_wasm_wrappers